fix: fix convert record batch (#521)

apache · Dec 29, 2022 · 354c585 · 354c585
1 parent 1660efe
commit 354c585
Showing 1 changed file with 23 additions and 22 deletions.
diff --git a/server/src/handlers/sql.rs b/server/src/handlers/sql.rs
@@ -208,32 +208,33 @@ fn convert_records(records: RecordBatchVec) -> ArrowResult<Response> {
         }));
     }
 
-    let record_batch = &records[0];
-    let num_cols = record_batch.num_columns();
-    let num_rows = record_batch.num_rows();
-    let schema = record_batch.schema();
-
-    let mut column_names = Vec::with_capacity(num_cols);
-    let mut column_data = Vec::with_capacity(num_rows);
-
-    for col_idx in 0..num_cols {
-        let column_schema = schema.column(col_idx).clone();
-        column_names.push(ResponseColumn {
-            name: column_schema.name,
-            data_type: column_schema.data_type,
-        });
-    }
+    let mut column_names = vec![];
+    let mut column_data = vec![];
 
-    for row_idx in 0..num_rows {
-        let mut row_data = Vec::with_capacity(num_cols);
-        for col_idx in 0..num_cols {
-            let column = record_batch.column(col_idx);
-            let column = column.datum(row_idx);
+    for record_batch in records {
+        let num_cols = record_batch.num_columns();
+        let num_rows = record_batch.num_rows();
+        let schema = record_batch.schema();
 
-            row_data.push(column);
+        for col_idx in 0..num_cols {
+            let column_schema = schema.column(col_idx).clone();
+            column_names.push(ResponseColumn {
+                name: column_schema.name,
+                data_type: column_schema.data_type,
+            });
         }
 
-        column_data.push(row_data);
+        for row_idx in 0..num_rows {
+            let mut row_data = Vec::with_capacity(num_cols);
+            for col_idx in 0..num_cols {
+                let column = record_batch.column(col_idx);
+                let column = column.datum(row_idx);
+
+                row_data.push(column);
+            }
+
+            column_data.push(row_data);
+        }
     }
 
     Ok(Response::Rows(ResponseRows {