chore: table display discard some unused rows

Cargo.toml

@@ -29,10 +29,10 @@ databend-driver-macros = { path = "macros", version = "0.26.2" }
 
 tokio-stream = "0.1"
 chrono = { version = "=0.4.39", default-features = false, features = ["clock"] }
-arrow = { version = "53.0" }
-arrow-array = { version = "53.0" }
-arrow-schema = { version = "53.0" }
-arrow-flight = { version = "53.0", features = ["flight-sql-experimental"] }
+arrow = { version = "54.0" }
+arrow-array = { version = "54.0" }
+arrow-schema = { version = "54.0" }
+arrow-flight = { version = "54.0", features = ["flight-sql-experimental"] }
 tonic = { version = "0.12", default-features = false, features = [
     "transport",
     "codegen",

bindings/nodejs/tests/binding.js

@@ -294,7 +294,7 @@ Then("Stream load and Select should be equal", async function () {
   ];
   const progress = await this.conn.streamLoad(`INSERT INTO test VALUES`, values);
   assert.equal(progress.writeRows, 3);
-  assert.equal(progress.writeBytes, 187);
+  assert.equal(progress.writeBytes, 193);
 
   const rows = await this.conn.queryIter("SELECT * FROM test");
   const ret = [];
@@ -312,7 +312,7 @@ Then("Stream load and Select should be equal", async function () {
 Then("Load file and Select should be equal", async function () {
   const progress = await this.conn.loadFile(`INSERT INTO test VALUES`, "tests/data/test.csv", { type: "CSV" });
   assert.equal(progress.writeRows, 3);
-  assert.equal(progress.writeBytes, 187);
+  assert.equal(progress.writeBytes, 193);
 
   const rows = await this.conn.queryIter("SELECT * FROM test");
   const ret = [];

bindings/python/tests/asyncio/steps/binding.py

@@ -166,7 +166,7 @@ async def _(context):
     ]
     progress = await context.conn.stream_load("INSERT INTO test VALUES", values)
     assert progress.write_rows == 3, f"progress.write_rows: {progress.write_rows}"
-    assert progress.write_bytes == 187, f"progress.write_bytes: {progress.write_bytes}"
+    assert progress.write_bytes == 193, f"progress.write_bytes: {progress.write_bytes}"
 
     rows = await context.conn.query_iter("SELECT * FROM test")
     ret = []
@@ -186,7 +186,7 @@ async def _(context):
         "INSERT INTO test VALUES", "tests/data/test.csv", {"type": "CSV"}
     )
     assert progress.write_rows == 3, f"progress.write_rows: {progress.write_rows}"
-    assert progress.write_bytes == 187, f"progress.write_bytes: {progress.write_bytes}"
+    assert progress.write_bytes == 193, f"progress.write_bytes: {progress.write_bytes}"
 
     rows = await context.conn.query_iter("SELECT * FROM test")
     ret = []

bindings/python/tests/blocking/steps/binding.py

@@ -152,7 +152,7 @@ def _(context):
     ]
     progress = context.conn.stream_load("INSERT INTO test VALUES", values)
     assert progress.write_rows == 3, f"progress.write_rows: {progress.write_rows}"
-    assert progress.write_bytes == 187, f"progress.write_bytes: {progress.write_bytes}"
+    assert progress.write_bytes == 193, f"progress.write_bytes: {progress.write_bytes}"
 
     rows = context.conn.query_iter("SELECT * FROM test")
     ret = []
@@ -172,7 +172,7 @@ def _(context):
         "INSERT INTO test VALUES", "tests/data/test.csv", {"type": "CSV"}
     )
     assert progress.write_rows == 3, f"progress.write_rows: {progress.write_rows}"
-    assert progress.write_bytes == 187, f"progress.write_bytes: {progress.write_bytes}"
+    assert progress.write_bytes == 193, f"progress.write_bytes: {progress.write_bytes}"
 
     rows = context.conn.query_iter("SELECT * FROM test")
     ret = []

cli/src/display.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::VecDeque;
 use std::env;
 use std::fmt::Write;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -51,7 +52,7 @@ pub struct FormatDisplay<'a> {
     replace_newline: bool,
     data: RowStatsIterator,
 
-    rows: usize,
+    rows_count: usize,
     progress: Option<ProgressBar>,
     start: Instant,
     stats: Option<ServerStats>,
@@ -73,7 +74,7 @@ impl<'a> FormatDisplay<'a> {
             kind: QueryKind::from(query),
             replace_newline,
             data,
-            rows: 0,
+            rows_count: 0,
             progress: None,
             start,
             stats: None,
@@ -139,16 +140,28 @@ impl FormatDisplay<'_> {
             let format_sql = highlight_query(&format_sql);
             println!("\n{}\n", format_sql);
         }
+
+        let max_display_rows_count = self.settings.max_display_rows / 2 + 2;
         let mut rows = Vec::new();
+        let mut bottom_rows = VecDeque::new();
         let mut error = None;
         while let Some(line) = self.data.next().await {
             if self.interrupted.load(Ordering::SeqCst) {
                 return Err(anyhow!(INTERRUPTED_MESSAGE));
             }
             match line {
                 Ok(RowWithStats::Row(row)) => {
-                    self.rows += 1;
-                    rows.push(row);
+                    if self.rows_count < max_display_rows_count {
+                        rows.push(row);
+                    } else {
+                        bottom_rows.push_back(row);
+                        // Since bendsql only displays the maximum number of rows,
+                        // we can discard some rows to avoid data occupying too much memory.
+                        if bottom_rows.len() > max_display_rows_count {
+                            bottom_rows.pop_front();
+                        }
+                    }
+                    self.rows_count += 1;
                 }
                 Ok(RowWithStats::Stats(ss)) => {
                     self.display_progress(&ss).await;
@@ -160,6 +173,11 @@ impl FormatDisplay<'_> {
                 }
             }
         }
+        // collect bottom rows
+        while let Some(row) = bottom_rows.pop_front() {
+            rows.push(row);
+        }
+
         if let Some(pb) = self.progress.take() {
             pb.finish_and_clear();
         }
@@ -204,6 +222,7 @@ impl FormatDisplay<'_> {
                         self.settings.max_display_rows,
                         self.settings.max_width,
                         self.settings.max_col_width,
+                        self.rows_count
                     )?
                 );
             }
@@ -217,7 +236,8 @@ impl FormatDisplay<'_> {
                         self.replace_newline,
                         self.settings.max_display_rows,
                         self.settings.max_width,
-                        self.settings.max_col_width
+                        self.settings.max_col_width,
+                        self.rows_count
                     )?
                 );
             }
@@ -242,7 +262,7 @@ impl FormatDisplay<'_> {
             }
             match line {
                 Ok(RowWithStats::Row(row)) => {
-                    self.rows += 1;
+                    self.rows_count += 1;
                     let record = row.into_iter().map(|v| v.to_string()).collect::<Vec<_>>();
                     wtr.write_record(record)?;
                 }
@@ -275,7 +295,7 @@ impl FormatDisplay<'_> {
             }
             match line {
                 Ok(RowWithStats::Row(row)) => {
-                    self.rows += 1;
+                    self.rows_count += 1;
                     let record = row.into_iter().map(|v| v.to_string()).collect::<Vec<_>>();
                     wtr.write_record(record)?;
                 }
@@ -298,7 +318,7 @@ impl FormatDisplay<'_> {
             }
             match line {
                 Ok(RowWithStats::Row(_)) => {
-                    self.rows += 1;
+                    self.rows_count += 1;
                 }
                 Ok(RowWithStats::Stats(ss)) => {
                     self.display_progress(&ss).await;
@@ -316,7 +336,7 @@ impl FormatDisplay<'_> {
         if let Some(err) = error {
             return Err(anyhow!(
                 "error happens after fetched {} rows: {}",
-                self.rows,
+                self.rows_count,
                 err
             ));
         }
@@ -332,10 +352,16 @@ impl FormatDisplay<'_> {
             stats.normalize();
 
             let (rows, mut rows_str, kind, total_rows, total_bytes) = match self.kind {
-                QueryKind::Graphical => (self.rows, "rows", "graphical", 0, 0),
-                QueryKind::Explain => (self.rows, "rows", "explain", 0, 0),
-                QueryKind::ShowCreate => (self.rows, "rows", "showcreate", 0, 0),
-                QueryKind::Query => (self.rows, "rows", "read", stats.read_rows, stats.read_bytes),
+                QueryKind::Graphical => (self.rows_count, "rows", "graphical", 0, 0),
+                QueryKind::Explain => (self.rows_count, "rows", "explain", 0, 0),
+                QueryKind::ShowCreate => (self.rows_count, "rows", "showcreate", 0, 0),
+                QueryKind::Query => (
+                    self.rows_count,
+                    "rows",
+                    "read",
+                    stats.read_rows,
+                    stats.read_bytes,
+                ),
                 QueryKind::Update | QueryKind::AlterUserPassword | QueryKind::GenData(_, _, _) => (
                     stats.write_rows,
                     "rows",
@@ -484,6 +510,7 @@ fn create_table(
     max_rows: usize,
     max_width: usize,
     max_col_width: usize,
+    rows_count: usize,
 ) -> Result<Table> {
     let mut table = Table::new();
     table
@@ -506,21 +533,20 @@ fn create_table(
         return Ok(table);
     }
 
-    let row_count: usize = results.len();
-    let mut rows_to_render = row_count.min(max_rows);
+    let value_rows_count: usize = results.len();
+    let mut rows_to_render = value_rows_count.min(max_rows);
     if !replace_newline {
-        rows_to_render = row_count;
-    } else if row_count <= max_rows + 3 {
+        rows_to_render = value_rows_count;
+    } else if value_rows_count <= max_rows + 3 {
         // hiding rows adds 3 extra rows
         // so hiding rows makes no sense if we are only slightly over the limit
         // if we are 1 row over the limit hiding rows will actually increase the number of lines we display!
         // in this case render all the rows
-        // 	rows_to_render = row_count;
-        rows_to_render = row_count;
+        rows_to_render = value_rows_count;
     }
 
-    let (top_rows, bottom_rows) = if rows_to_render == row_count {
-        (row_count, 0usize)
+    let (top_rows, bottom_rows) = if rows_to_render == value_rows_count {
+        (value_rows_count, 0usize)
     } else {
         let top_rows = rows_to_render / 2 + (rows_to_render % 2 != 0) as usize;
         (top_rows, rows_to_render - top_rows)
@@ -547,7 +573,7 @@ fn create_table(
         });
 
     if bottom_rows != 0 {
-        for row in results.iter().skip(row_count - bottom_rows) {
+        for row in results.iter().skip(value_rows_count - bottom_rows) {
             let values = row.values();
             let mut v = vec![];
             for value in values {
@@ -597,9 +623,9 @@ fn create_table(
             table.add_row(cells);
         }
 
-        let row_count_str = format!("{} rows", row_count);
+        let rows_count_str = format!("{} rows", rows_count);
         let show_count_str = format!("({} shown)", top_rows + bottom_rows);
-        table.add_row(vec![Cell::new(row_count_str).set_alignment(aligns[0])]);
+        table.add_row(vec![Cell::new(rows_count_str).set_alignment(aligns[0])]);
         table.add_row(vec![Cell::new(show_count_str).set_alignment(aligns[0])]);
     }
 

cli/src/session.rs

@@ -383,11 +383,12 @@ impl Session {
                 },
             }
         }
+        // save history first to avoid loss data.
+        let _ = rl.save_history(&get_history_path());
         if let Err(e) = self.conn.close().await {
             println!("got error when closing session: {}", e);
         }
         println!("Bye~");
-        let _ = rl.save_history(&get_history_path());
     }
 
     pub async fn handle_reader<R: BufRead>(&mut self, r: R) -> Result<()> {

sql/src/raw_rows.rs

@@ -63,9 +63,8 @@ impl TryFrom<(SchemaRef, Vec<Option<String>>)> for RawRow {
     type Error = Error;
 
     fn try_from((schema, data): (SchemaRef, Vec<Option<String>>)) -> Result<Self> {
-        let mut values: Vec<Value> = Vec::new();
-        for (i, field) in schema.fields().iter().enumerate() {
-            let val: Option<&str> = data.get(i).and_then(|v| v.as_deref());
+        let mut values: Vec<Value> = Vec::with_capacity(data.len());
+        for (field, val) in schema.fields().iter().zip(data.clone().into_iter()) {
             values.push(Value::try_from((&field.data_type, val))?);
         }
 

sql/src/rows.rs

@@ -139,9 +139,8 @@ impl TryFrom<(SchemaRef, Vec<Option<String>>)> for Row {
     type Error = Error;
 
     fn try_from((schema, data): (SchemaRef, Vec<Option<String>>)) -> Result<Self> {
-        let mut values: Vec<Value> = Vec::new();
-        for (i, field) in schema.fields().iter().enumerate() {
-            let val: Option<&str> = data.get(i).and_then(|v| v.as_deref());
+        let mut values: Vec<Value> = Vec::with_capacity(data.len());
+        for (field, val) in schema.fields().iter().zip(data.into_iter()) {
             values.push(Value::try_from((&field.data_type, val))?);
         }
         Ok(Self::new(schema, values))

sql/src/value.rs

@@ -154,10 +154,10 @@ impl Value {
     }
 }
 
-impl TryFrom<(&DataType, Option<&str>)> for Value {
+impl TryFrom<(&DataType, Option<String>)> for Value {
     type Error = Error;
 
-    fn try_from((t, v): (&DataType, Option<&str>)) -> Result<Self> {
+    fn try_from((t, v): (&DataType, Option<String>)) -> Result<Self> {
         match v {
             Some(v) => Self::try_from((t, v)),
             None => match t {
@@ -171,17 +171,17 @@ impl TryFrom<(&DataType, Option<&str>)> for Value {
     }
 }
 
-impl TryFrom<(&DataType, &str)> for Value {
+impl TryFrom<(&DataType, String)> for Value {
     type Error = Error;
 
-    fn try_from((t, v): (&DataType, &str)) -> Result<Self> {
+    fn try_from((t, v): (&DataType, String)) -> Result<Self> {
         match t {
             DataType::Null => Ok(Self::Null),
             DataType::EmptyArray => Ok(Self::EmptyArray),
             DataType::EmptyMap => Ok(Self::EmptyMap),
             DataType::Boolean => Ok(Self::Boolean(v == "1")),
             DataType::Binary => Ok(Self::Binary(hex::decode(v)?)),
-            DataType::String => Ok(Self::String(v.to_string())),
+            DataType::String => Ok(Self::String(v)),
             DataType::Number(NumberDataType::Int8) => {
                 Ok(Self::Number(NumberValue::Int8(v.parse()?)))
             }
@@ -213,28 +213,29 @@ impl TryFrom<(&DataType, &str)> for Value {
                 Ok(Self::Number(NumberValue::Float64(v.parse()?)))
             }
             DataType::Decimal(DecimalDataType::Decimal128(size)) => {
-                let d = parse_decimal(v, *size)?;
+                let d = parse_decimal(v.as_str(), *size)?;
                 Ok(Self::Number(d))
             }
             DataType::Decimal(DecimalDataType::Decimal256(size)) => {
-                let d = parse_decimal(v, *size)?;
+                let d = parse_decimal(v.as_str(), *size)?;
                 Ok(Self::Number(d))
             }
             DataType::Timestamp => Ok(Self::Timestamp(
-                NaiveDateTime::parse_from_str(v, "%Y-%m-%d %H:%M:%S%.6f")?
+                NaiveDateTime::parse_from_str(v.as_str(), "%Y-%m-%d %H:%M:%S%.6f")?
                     .and_utc()
                     .timestamp_micros(),
             )),
             DataType::Date => Ok(Self::Date(
-                NaiveDate::parse_from_str(v, "%Y-%m-%d")?.num_days_from_ce() - DAYS_FROM_CE,
+                NaiveDate::parse_from_str(v.as_str(), "%Y-%m-%d")?.num_days_from_ce()
+                    - DAYS_FROM_CE,
             )),
-            DataType::Bitmap => Ok(Self::Bitmap(v.to_string())),
-            DataType::Variant => Ok(Self::Variant(v.to_string())),
-            DataType::Geometry => Ok(Self::Geometry(v.to_string())),
-            DataType::Geography => Ok(Self::Geography(v.to_string())),
-            DataType::Interval => Ok(Self::Interval(v.to_string())),
+            DataType::Bitmap => Ok(Self::Bitmap(v)),
+            DataType::Variant => Ok(Self::Variant(v)),
+            DataType::Geometry => Ok(Self::Geometry(v)),
+            DataType::Geography => Ok(Self::Geography(v)),
+            DataType::Interval => Ok(Self::Interval(v)),
             DataType::Array(_) | DataType::Map(_) | DataType::Tuple(_) => {
-                let mut reader = Cursor::new(v);
+                let mut reader = Cursor::new(v.as_str());
                 let decoder = ValueDecoder {};
                 decoder.read_field(t, &mut reader)
             }