chore(cubestore): Upgrade DF: Use DF 46.0.1

Author: srh

rust/cubestore/Cargo.lock

@@ -18,7 +18,7 @@ base64 = "0.13.0"
 bumpalo = "3.6.1"
 tokio = { version = "1", features = ["full", "rt"] }
 warp = { version = "0.3.6" }
-sqlparser = { git = "https://github.com/cube-js/sqlparser-rs.git", branch = "cube-42.2.0" }
+sqlparser = { git = "https://github.com/cube-js/sqlparser-rs.git", branch = "cube-46.0.1" }
 serde_derive = "1.0.115"
 serde = "1.0.115"
 serde_repr = "0.1"
@@ -29,9 +29,10 @@ cubezetasketch = { path = "../cubezetasketch" }
 cubedatasketches = { path = "../cubedatasketches" }
 cubeshared = { path = "../../cubeshared" }
 cuberpc = { path = "../cuberpc" }
-datafusion = { git = "https://github.com/cube-js/arrow-datafusion", branch = "cube-42.2.0", features = ["serde"] }
-datafusion-proto = { git = "https://github.com/cube-js/arrow-datafusion", branch = "cube-42.2.0" }
-datafusion-proto-common = { git = "https://github.com/cube-js/arrow-datafusion", branch = "cube-42.2.0" }
+datafusion = { git = "https://github.com/cube-js/arrow-datafusion", branch = "cube-46.0.1", features = ["serde"] }
+datafusion-datasource = { git = "https://github.com/cube-js/arrow-datafusion", branch = "cube-46.0.1" }
+datafusion-proto = { git = "https://github.com/cube-js/arrow-datafusion", branch = "cube-46.0.1" }
+datafusion-proto-common = { git = "https://github.com/cube-js/arrow-datafusion", branch = "cube-46.0.1" }
 csv = "1.1.3"
 bytes = "1.6.0"
 serde_json = "1.0.56"

rust/cubestore/cubestore/Cargo.toml

@@ -341,7 +341,7 @@ impl DataFrameValue<String> for Option<Vec<AggregateFunction>> {
     }
 }
 
-#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, DeepSizeOf)]
+#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd, DeepSizeOf)]
 pub enum HllFlavour {
     Airlift,      // Compatible with Presto, Athena, etc.
     Snowflake,    // Same storage as Airlift, imports from Snowflake JSON.
@@ -369,7 +369,7 @@ pub fn is_valid_plain_binary_hll(data: &[u8], f: HllFlavour) -> Result<(), CubeE
     return Ok(());
 }
 
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, DeepSizeOf)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd, DeepSizeOf)]
 pub enum ColumnType {
     String,
     Int,
@@ -547,7 +547,7 @@ impl From<&Column> for types::Type {
     }
 }
 
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, DeepSizeOf)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd, DeepSizeOf)]
 pub struct Column {
     name: String,
     column_type: ColumnType,
@@ -611,7 +611,7 @@ impl fmt::Display for Column {
     }
 }
 
-#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub enum ImportFormat {
     CSV,
     CSVNoHeader,
@@ -624,22 +624,22 @@ pub enum ImportFormat {
 }
 
 data_frame_from! {
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub struct Schema {
     name: String
 }
 }
 
 impl RocksEntity for Schema {}
 
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub enum IndexType {
     Regular = 1,
     Aggregate = 2,
 }
 
 data_frame_from! {
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub struct Index {
     name: String,
     table_id: u64,
@@ -656,7 +656,7 @@ pub struct Index {
 
 impl RocksEntity for Index {}
 
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub enum AggregateFunction {
     SUM = 1,
     MAX = 2,

rust/cubestore/cubestore/src/metastore/mod.rs

@@ -598,7 +598,7 @@ impl WriteBatchIterator for WriteBatchContainer {
     }
 }
 
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub struct IdRow<T: Clone> {
     pub(crate) id: u64,
     pub(crate) row: T,

rust/cubestore/cubestore/src/metastore/rocks_store.rs

@@ -23,7 +23,7 @@ use serde::{Deserialize, Deserializer, Serialize};
 use std::io::Write;
 use std::sync::Arc;
 
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub struct AggregateColumnIndex {
     index: u64,
     function: AggregateFunction,
@@ -114,7 +114,7 @@ impl core::fmt::Display for AggregateColumn {
     }
 }
 
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub enum StreamOffset {
     Earliest = 1,
     Latest = 2,
@@ -129,7 +129,7 @@ impl DataFrameValue<String> for Option<StreamOffset> {
 }
 
 data_frame_from! {
-#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash)]
+#[derive(Clone, Serialize, Deserialize, Debug, Eq, PartialEq, Hash, PartialOrd)]
 pub struct Table {
     table_name: String,
     schema_id: u64,
@@ -172,7 +172,7 @@ pub struct Table {
 
 impl RocksEntity for Table {}
 
-#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash, PartialOrd)]
 pub struct TablePath {
     pub table: IdRow<Table>,
     pub schema: Arc<IdRow<Schema>>,

rust/cubestore/cubestore/src/metastore/table.rs

@@ -47,15 +47,13 @@ impl OptimizerRule for FlattenUnion {
             | LogicalPlan::Values(_)
             | LogicalPlan::Analyze(_)
             | LogicalPlan::Distinct(_)
-            | LogicalPlan::Prepare(_)
             // | LogicalPlan::Execute(_)
             | LogicalPlan::Dml(_)
             | LogicalPlan::Ddl(_)
             | LogicalPlan::Copy(_)
             | LogicalPlan::DescribeTable(_)
             | LogicalPlan::Unnest(_)
             | LogicalPlan::RecursiveQuery(_)
-            | LogicalPlan::CrossJoin(_)
             => {
                 // apply the optimization to all inputs of the plan
                 let inputs = plan.inputs();

rust/cubestore/cubestore/src/queryplanner/flatten_union.rs

@@ -3,6 +3,8 @@ pub mod optimizations;
 pub mod panic;
 mod partition_filter;
 mod planning;
+use datafusion::logical_expr::planner::ExprPlanner;
+use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 // use datafusion::physical_plan::parquet::MetadataCacheFactory;
 pub use planning::PlanningMeta;
 mod check_memory;
@@ -81,10 +83,11 @@ use datafusion::logical_expr::{
     TableSource, WindowUDF,
 };
 use datafusion::physical_expr::EquivalenceProperties;
-use datafusion::physical_plan::memory::MemoryExec;
+// TODO upgrade DF
+// use datafusion::physical_plan::memory::MemoryExec;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
-    collect, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, Partitioning,
+    collect, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning,
     PlanProperties, SendableRecordBatchStream,
 };
 use datafusion::prelude::{SessionConfig, SessionContext};
@@ -288,6 +291,7 @@ struct MetaStoreSchemaProvider {
     inline_tables: InlineTables,
     cache: Arc<SqlResultCache>,
     config_options: ConfigOptions,
+    expr_planners: Vec<Arc<dyn ExprPlanner>>,  // session_state.expr_planners clone
     session_state: Arc<SessionState>,
 }
 
@@ -333,6 +337,7 @@ impl MetaStoreSchemaProvider {
             cache,
             inline_tables: (*inline_tables).clone(),
             config_options: ConfigOptions::new(),
+            expr_planners: datafusion::execution::FunctionRegistry::expr_planners(session_state.as_ref()),
             session_state,
         }
     }
@@ -572,6 +577,11 @@ impl ContextProvider for MetaStoreSchemaProvider {
             .cloned()
             .collect()
     }
+
+    // We implement this for count(*) replacement.
+    fn get_expr_planners(&self) -> &[Arc<dyn datafusion::logical_expr::planner::ExprPlanner>] {
+        self.expr_planners.as_slice()
+    }
 }
 
 /// Enables our options used with `SqlToRel`.  Sets `enable_ident_normalization` to false.  See also
@@ -760,7 +770,8 @@ impl TableProvider for InfoSchemaTableProvider {
             properties: PlanProperties::new(
                 EquivalenceProperties::new(schema),
                 Partitioning::UnknownPartitioning(1),
-                ExecutionMode::Bounded,
+                EmissionType::Both,  // TODO upgrade DF:  Both is safe choice
+                Boundedness::Bounded,
             ),
         };
         Ok(Arc::new(exec))

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

@@ -2,9 +2,10 @@ use crate::queryplanner::check_memory::CheckMemoryExec;
 use crate::queryplanner::query_executor::ClusterSendExec;
 use crate::util::memory::MemoryHandler;
 use datafusion::datasource::physical_plan::ParquetExec;
+use datafusion::datasource::source::DataSourceExec;
 use datafusion::error::DataFusionError;
-use datafusion::physical_plan::memory::MemoryExec;
 use datafusion::physical_plan::ExecutionPlan;
+use datafusion_datasource::memory::MemoryExec;
 use std::sync::Arc;
 
 /// Add `CheckMemoryExec` behind some nodes.
@@ -13,7 +14,8 @@ pub fn add_check_memory_exec(
     mem_handler: Arc<dyn MemoryHandler>,
 ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
     let p_any = p.as_any();
-    if p_any.is::<ParquetExec>() || p_any.is::<MemoryExec>() || p_any.is::<ClusterSendExec>() {
+    // TODO upgrade DF: Do we use ParquetExec?  Or just DataSourceExec?  It's fine to have both here.
+    if p_any.is::<DataSourceExec>() || p_any.is::<ParquetExec>() || p_any.is::<MemoryExec>() || p_any.is::<ClusterSendExec>() {
         let memory_check = Arc::new(CheckMemoryExec::new(p, mem_handler.clone()));
         Ok(memory_check)
     } else {

rust/cubestore/cubestore/src/queryplanner/optimizations/check_memory.rs

@@ -4,6 +4,7 @@ use crate::queryplanner::query_executor::ClusterSendExec;
 use crate::queryplanner::tail_limit::TailLimitExec;
 use crate::queryplanner::topk::AggregateTopKExec;
 use datafusion::error::DataFusionError;
+use datafusion::physical_expr::LexOrdering;
 use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode};
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::limit::GlobalLimitExec;
@@ -122,11 +123,11 @@ pub fn ensure_partition_merge_helper(
                 .children()
                 .into_iter()
                 .map(|c| -> Arc<dyn ExecutionPlan> {
-                    Arc::new(SortPreservingMergeExec::new(ordering.clone(), c.clone()))
+                    Arc::new(SortPreservingMergeExec::new(LexOrdering::new(ordering.clone()), c.clone()))
                 })
                 .collect();
             let new_plan = p.clone().with_new_children(merged_children)?;
-            Arc::new(SortPreservingMergeExec::new(ordering, new_plan))
+            Arc::new(SortPreservingMergeExec::new(LexOrdering::new(ordering), new_plan))
         } else {
             let merged_children = p
                 .children()

rust/cubestore/cubestore/src/queryplanner/optimizations/distributed_partial_aggregate.rs

@@ -107,6 +107,7 @@ impl QueryPlanner for CubeQueryPlanner {
     }
 }
 
+#[derive(Debug)]
 pub struct PreOptimizeRule {
     memory_handler: Arc<dyn MemoryHandler>,
     data_loaded_size: Option<Arc<DataLoadedSize>>,