fix(cubestore): Reduce memory usage while converting to DataFrame (#8599)

Author: ovr

rust/cubestore/cubestore/src/queryplanner/mod.rs

@@ -40,7 +40,7 @@ use crate::queryplanner::info_schema::{
 use crate::queryplanner::now::MaterializeNow;
 use crate::queryplanner::planning::{choose_index_ext, ClusterSendNode};
 use crate::queryplanner::query_executor::{
-    batch_to_dataframe, ClusterSendExec, InlineTableProvider,
+    batches_to_dataframe, ClusterSendExec, InlineTableProvider,
 };
 use crate::queryplanner::serialized_plan::SerializedPlan;
 use crate::queryplanner::topk::ClusterAggregateTopK;
@@ -168,7 +168,7 @@ impl QueryPlanner for QueryPlannerImpl {
         let execution_time = execution_time.elapsed()?;
         app_metrics::META_QUERY_TIME_MS.report(execution_time.as_millis() as i64);
         debug!("Meta query data processing time: {:?}", execution_time,);
-        let data_frame = cube_ext::spawn_blocking(move || batch_to_dataframe(&results)).await??;
+        let data_frame = cube_ext::spawn_blocking(move || batches_to_dataframe(results)).await??;
         Ok(data_frame)
     }
 }

rust/cubestore/cubestore/src/queryplanner/query_executor.rs

@@ -1409,25 +1409,26 @@ macro_rules! convert_array {
     }};
 }
 
-pub fn batch_to_dataframe(batches: &Vec<RecordBatch>) -> Result<DataFrame, CubeError> {
+pub fn batches_to_dataframe(batches: Vec<RecordBatch>) -> Result<DataFrame, CubeError> {
     let mut cols = vec![];
     let mut all_rows = vec![];
 
-    for batch in batches.iter() {
+    for batch in batches.into_iter() {
         if cols.len() == 0 {
-            let schema = batch.schema().clone();
-            for (i, field) in schema.fields().iter().enumerate() {
+            for (i, field) in batch.schema().fields().iter().enumerate() {
                 cols.push(Column::new(
                     field.name().clone(),
                     arrow_to_column_type(field.data_type().clone())?,
                     i,
                 ));
             }
         }
+
         if batch.num_rows() == 0 {
             continue;
         }
-        let mut rows = vec![];
+
+        let mut rows = Vec::with_capacity(batch.num_rows());
 
         for _ in 0..batch.num_rows() {
             rows.push(Row::new(Vec::with_capacity(batch.num_columns())));
@@ -1769,7 +1770,7 @@ mod tests {
             Field::new("int32", DataType::Int32, true),
             Field::new("str32", DataType::Utf8, true),
         ]));
-        let result = batch_to_dataframe(&vec![RecordBatch::try_new(
+        let result = batches_to_dataframe(vec![RecordBatch::try_new(
             schema,
             vec![
                 Arc::new(UInt32Array::from_iter(vec![Some(1), None])) as ArrayRef,

rust/cubestore/cubestore/src/sql/mod.rs

@@ -49,7 +49,7 @@ use crate::metastore::{
 };
 use crate::queryplanner::panic::PanicWorkerNode;
 use crate::queryplanner::pretty_printers::{pp_phys_plan, pp_plan};
-use crate::queryplanner::query_executor::{batch_to_dataframe, ClusterSendExec, QueryExecutor};
+use crate::queryplanner::query_executor::{batches_to_dataframe, ClusterSendExec, QueryExecutor};
 use crate::queryplanner::serialized_plan::{RowFilter, SerializedPlan};
 use crate::queryplanner::{PlanningMeta, QueryPlan, QueryPlanner};
 use crate::remotefs::RemoteFs;
@@ -1072,7 +1072,7 @@ impl SqlService for SqlServiceImpl {
                                     }
                                     Ok(cube_ext::spawn_blocking(
                                         move || -> Result<DataFrame, CubeError> {
-                                            let df = batch_to_dataframe(&records)?;
+                                            let df = batches_to_dataframe(records)?;
                                             Ok(df)
                                         },
                                     )

rust/cubestore/cubestore/src/streaming/kafka.rs

@@ -402,7 +402,7 @@ impl StreamingSource for KafkaStreamingSource {
 mod tests {
     use super::*;
     use crate::metastore::{Column, ColumnType};
-    use crate::queryplanner::query_executor::batch_to_dataframe;
+    use crate::queryplanner::query_executor::batches_to_dataframe;
     use crate::sql::MySqlDialectWithBackTicks;
     use crate::streaming::topic_table_provider::TopicTableProvider;
     use arrow::array::StringArray;
@@ -432,7 +432,7 @@ mod tests {
         let phys_plan = plan_ctx.create_physical_plan(&logical_plan).unwrap();
 
         let batches = collect(phys_plan).await.unwrap();
-        let res = batch_to_dataframe(&batches).unwrap();
+        let res = batches_to_dataframe(batches).unwrap();
         res.get_rows()[0].values()[0].clone()
     }
 
@@ -462,7 +462,7 @@ mod tests {
         let phys_plan = phys_plan.with_new_children(vec![inp]).unwrap();
 
         let batches = collect(phys_plan).await.unwrap();
-        let res = batch_to_dataframe(&batches).unwrap();
+        let res = batches_to_dataframe(batches).unwrap();
         res.get_rows().to_vec()
     }