Upgrade Datafusion to 52 and fix the schema errors

Cargo.toml

@@ -126,21 +126,14 @@ prost-build = "0.13.1"
 lazy_static = "1.4.0"
 pest = "2.7.8"
 pest_derive = "2.7.8"
-sqlparser = "0.43.1"
+sqlparser = "0.47"
 
-datafusion = { version = "36" }
+datafusion = { version = "40" }
 futures = "0.3"
-arrow = { version = "50" }
-arrow-buffer = { version = "50" }
-arrow-schema = { version = "50" }
-arrow-array = { version = "50" }
+arrow = { version = "52" }
+arrow-buffer = { version = "52" }
+arrow-schema = { version = "52" }
+arrow-data = { version = "52" }
+arrow-array = { version = "52" }
 
 moka = { version = "0.12.7", features = ["sync"] }
-
-# Make sure that transitive dependencies stick to disk_graph 50
-[patch.crates-io]
-arrow = { git = "https://github.com/apache/arrow-rs.git", tag = "50.0.0" }
-arrow-buffer = { git = "https://github.com/apache/arrow-rs.git", tag = "50.0.0" }
-arrow-schema = { git = "https://github.com/apache/arrow-rs.git", tag = "50.0.0" }
-arrow-data = { git = "https://github.com/apache/arrow-rs.git", tag = "50.0.0" }
-arrow-array = { git = "https://github.com/apache/arrow-rs.git", tag = "50.0.0" }

raphtory-cypher/src/executor/table_provider/edge.rs

@@ -22,7 +22,7 @@ use datafusion::{
     physical_expr::PhysicalSortExpr,
     physical_plan::{
         metrics::MetricsSet, stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType,
-        ExecutionPlan,
+        ExecutionPlan, PlanProperties,
     },
     physical_planner::create_physical_sort_expr,
 };
@@ -32,6 +32,8 @@ use raphtory::disk_graph::DiskGraphStorage;
 
 use crate::executor::{arrow2_to_arrow_buf, ExecError};
 
+use super::plan_properties;
+
 // use super::plan_properties;
 
 pub struct EdgeListTableProvider {
@@ -146,11 +148,10 @@ impl TableProvider for EdgeListTableProvider {
         _limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
         let schema = projection
-            .as_ref()
-            .map(|proj| Arc::new(self.schema().project(proj).expect("failed projection")))
-            .unwrap_or_else(|| self.schema().clone());
+            .map(|proj| self.schema().project(proj).map(Arc::new))
+            .unwrap_or_else(|| Ok(self.schema().clone()))?;
 
-        // let plan_properties = plan_properties(schema.clone(), self.num_partitions);
+        let plan_properties = plan_properties(schema.clone(), self.num_partitions);
         Ok(Arc::new(EdgeListExecPlan {
             layer_id: self.layer_id,
             layer_name: self.layer_name.clone(),
@@ -159,7 +160,7 @@ impl TableProvider for EdgeListTableProvider {
             num_partitions: self.num_partitions,
             row_count: self.row_count,
             sorted_by: self.sorted_by.clone(),
-            // props: plan_properties,
+            props: plan_properties,
             projection: projection.map(|proj| Arc::from(proj.as_slice())),
         }))
     }
@@ -173,7 +174,7 @@ struct EdgeListExecPlan {
     num_partitions: usize,
     row_count: usize,
     sorted_by: Vec<PhysicalSortExpr>,
-    // props: PlanProperties,
+    props: PlanProperties,
     projection: Option<Arc<[usize]>>,
 }
 
@@ -349,19 +350,16 @@ impl DisplayAs for EdgeListExecPlan {
 
 #[async_trait]
 impl ExecutionPlan for EdgeListExecPlan {
+    fn name(&self) -> &str {
+        "EdgeListExecPlan"
+    }
+
     fn as_any(&self) -> &dyn Any {
         self
     }
 
-    // fn properties(&self) -> &PlanProperties {
-    //     &self.props
-    // }
-
-    fn output_partitioning(&self) -> datafusion::physical_expr::Partitioning {
-        datafusion::physical_expr::Partitioning::UnknownPartitioning(self.num_partitions)
-    }
-    fn output_ordering(&self) -> Option<&[datafusion::physical_expr::PhysicalSortExpr]> {
-        Some(&self.sorted_by)
+    fn properties(&self) -> &PlanProperties {
+        &self.props
     }
 
     fn schema(&self) -> SchemaRef {
@@ -372,7 +370,7 @@ impl ExecutionPlan for EdgeListExecPlan {
         vec![true; self.children().len()]
     }
 
-    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         vec![]
     }
 
@@ -388,7 +386,7 @@ impl ExecutionPlan for EdgeListExecPlan {
         target_partitions: usize,
         _config: &ConfigOptions,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>, DataFusionError> {
-        // let plan_properties = plan_properties(self.schema.clone(), target_partitions);
+        let plan_properties = plan_properties(self.schema.clone(), target_partitions);
         Ok(Some(Arc::new(EdgeListExecPlan {
             layer_id: self.layer_id,
             layer_name: self.layer_name.clone(),
@@ -397,7 +395,7 @@ impl ExecutionPlan for EdgeListExecPlan {
             num_partitions: target_partitions,
             row_count: self.row_count,
             sorted_by: self.sorted_by.clone(),
-            // props: plan_properties,
+            props: plan_properties,
             projection: self.projection.clone(),
         })))
     }

raphtory-cypher/src/executor/table_provider/mod.rs

@@ -1,23 +1,11 @@
-// use arrow_schema::SchemaRef;
-// use datafusion::{
-//     physical_expr::EquivalenceProperties,
-//     // physical_plan::{ExecutionMode, Partitioning, PlanProperties},
-//     physical_plan::Partitioning,
-// };
+use arrow_schema::SchemaRef;
+use datafusion::{physical_expr::EquivalenceProperties, physical_plan::PlanProperties};
 
 pub mod edge;
 pub mod node;
-// FIXME this error shows up in datafusion 37 raised https://github.com/apache/datafusion/issues/10421
-// called `Result::unwrap()` on an `Err` value: Context("EnforceDistribution", Internal("PhysicalOptimizer rule 'EnforceDistribution' failed, due to generate a different schema,
-// schema: Schema { fields: [Field { name: \"name\", data_type: LargeUtf8, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: \"name\", data_type: LargeUtf8, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: \"name\", data_type: LargeUtf8, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} },
-// schema: Schema { fields: [Field { name: \"name\", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: \"name\", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: \"name\", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }"))
-
-// pub fn plan_properties(schema: SchemaRef, target_partitions: usize) -> PlanProperties {
-//     let eq_properties = EquivalenceProperties::new(schema.clone());
-//     let plan_properties = PlanProperties::new(
-//         eq_properties,
-//         Partitioning::UnknownPartitioning(target_partitions),
-//         ExecutionMode::Bounded,
-//     );
-//     plan_properties
-// }
+pub fn plan_properties(schema: SchemaRef, num_partitions: usize) -> PlanProperties {
+    let eq_properties = EquivalenceProperties::new(schema);
+    let partitioning = datafusion::physical_plan::Partitioning::UnknownPartitioning(num_partitions);
+    let execution_mode = datafusion::physical_plan::ExecutionMode::Bounded;
+    PlanProperties::new(eq_properties, partitioning, execution_mode)
+}

raphtory-cypher/src/executor/table_provider/node.rs

@@ -17,7 +17,7 @@ use datafusion::{
     logical_expr::Expr,
     physical_plan::{
         metrics::MetricsSet, stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType,
-        ExecutionPlan,
+        ExecutionPlan, PlanProperties,
     },
 };
 use futures::Stream;
@@ -28,6 +28,9 @@ use raphtory::{
 };
 use std::{any::Any, fmt::Formatter, sync::Arc};
 
+use super::plan_properties;
+
+// FIXME: review this file, some of the assuptions and mapping between partitions and chunk sizes are not correct
 pub struct NodeTableProvider {
     graph: DiskGraphStorage,
     schema: SchemaRef,
@@ -116,18 +119,17 @@ impl TableProvider for NodeTableProvider {
         _limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
         let schema = projection
-            .as_ref()
-            .map(|proj| Arc::new(self.schema().project(proj).expect("failed projection")))
-            .unwrap_or_else(|| self.schema().clone());
+            .map(|proj| self.schema().project(proj).map(Arc::new))
+            .unwrap_or_else(|| Ok(self.schema().clone()))?;
 
-        // let plan_properties = plan_properties(self.schema.clone(), self.num_partitions);
+        let plan_properties = plan_properties(schema.clone(), self.num_partitions);
 
         Ok(Arc::new(NodeScanExecPlan {
             graph: self.graph.clone(),
             schema,
             num_partitions: self.num_partitions,
             chunk_size: self.chunk_size,
-            // props: plan_properties,
+            props: plan_properties,
             projection: projection.map(|proj| Arc::from(proj.as_slice())),
         }))
     }
@@ -154,12 +156,15 @@ async fn produce_record_batch(
     let start = chunk_id * chunk_size;
     let end = (chunk_id + 1) * chunk_size;
 
+    let n = chunk.values()[0].len();
+    let iter = (start as u64..end as u64).take(n);
     let id = Arc::new(PrimitiveArray::<UInt64Type>::new(
-        ScalarBuffer::from_iter((start as u64..end as u64).take(chunk.values()[0].len())),
+        ScalarBuffer::from_iter(iter),
         None,
     ));
 
-    let arr_gid = graph.global_ordering().sliced(start, end - start);
+    let length = (end - start).min(graph.global_ordering().len());
+    let arr_gid = graph.global_ordering().sliced(start, length);
     let gid_data = to_data(arr_gid.as_ref());
     let gid = make_array(gid_data);
 
@@ -190,7 +195,7 @@ struct NodeScanExecPlan {
     schema: SchemaRef,
     num_partitions: usize,
     chunk_size: usize,
-    // props: PlanProperties,
+    props: PlanProperties,
     projection: Option<Arc<[usize]>>,
 }
 
@@ -231,19 +236,16 @@ impl DisplayAs for NodeScanExecPlan {
 
 #[async_trait]
 impl ExecutionPlan for NodeScanExecPlan {
+    fn name(&self) -> &str {
+        "NodeScanExecPlan"
+    }
+
     fn as_any(&self) -> &dyn Any {
         self
     }
 
-    // fn properties(&self) -> &PlanProperties {
-    //     &self.props
-    // }
-
-    fn output_partitioning(&self) -> datafusion::physical_expr::Partitioning {
-        datafusion::physical_expr::Partitioning::UnknownPartitioning(self.num_partitions)
-    }
-    fn output_ordering(&self) -> Option<&[datafusion::physical_expr::PhysicalSortExpr]> {
-        None
+    fn properties(&self) -> &PlanProperties {
+        &self.props
     }
 
     fn schema(&self) -> SchemaRef {
@@ -254,7 +256,7 @@ impl ExecutionPlan for NodeScanExecPlan {
         vec![true; self.children().len()]
     }
 
-    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
         vec![]
     }
 

raphtory-cypher/src/hop/execution.rs

@@ -8,7 +8,10 @@ use std::{
     task::{Context, Poll},
 };
 
-use crate::arrow2::{offset::Offset, types::NativeType};
+use crate::{
+    arrow2::{offset::Offset, types::NativeType},
+    executor::table_provider::plan_properties,
+};
 // use disk_graph::compute::take_record_batch;
 use arrow_array::{
     builder::{
@@ -26,11 +29,7 @@ use datafusion::{
     error::DataFusionError,
     execution::{RecordBatchStream, TaskContext},
     physical_plan::{
-        DisplayAs,
-        DisplayFormatType,
-        Distribution,
-        ExecutionPlan, //ExecutionPlanProperties,
-        // PlanProperties,
+        DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, PlanProperties,
         SendableRecordBatchStream,
     },
 };
@@ -60,7 +59,7 @@ pub struct HopExec {
     right_schema: DFSchemaRef,
 
     output_schema: SchemaRef,
-
+    props: PlanProperties,
     right_proj: Option<Vec<usize>>,
 }
 
@@ -85,7 +84,11 @@ impl HopExec {
 
         let input_col = find_last_input_col(hop, &input);
 
-        let out_schema: Schema = hop.out_schema.as_ref().into();
+        let out_schema: Arc<Schema> = Arc::new(hop.out_schema.as_ref().into());
+        let props = plan_properties(
+            out_schema.clone(),
+            input.properties().output_partitioning().partition_count(),
+        );
 
         Self {
             graph,
@@ -95,8 +98,8 @@ impl HopExec {
             right_schema: hop.right_schema.clone(),
             layers: hop.right_layers.clone(),
 
-            output_schema: Arc::new(out_schema),
-
+            output_schema: out_schema,
+            props,
             right_proj: hop.right_proj.clone(),
         }
     }
@@ -110,21 +113,22 @@ impl DisplayAs for HopExec {
 
 #[async_trait]
 impl ExecutionPlan for HopExec {
+    fn name(&self) -> &str {
+        "HopExec"
+    }
+
     /// Return a reference to Any that can be used for downcasting
     fn as_any(&self) -> &dyn Any {
         self
     }
 
+    fn properties(&self) -> &PlanProperties {
+        &self.props
+    }
+
     fn schema(&self) -> SchemaRef {
         self.output_schema.clone()
     }
-    fn output_partitioning(&self) -> Partitioning {
-        self.input.output_partitioning()
-    }
-    fn output_ordering(&self) -> Option<&[datafusion::physical_expr::PhysicalSortExpr]> {
-        self.input.output_ordering()
-    }
-
     fn required_input_distribution(&self) -> Vec<Distribution> {
         vec![Distribution::UnspecifiedDistribution]
     }
@@ -133,8 +137,8 @@ impl ExecutionPlan for HopExec {
         vec![true]
     }
 
-    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
-        vec![self.input.clone()]
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
     }
 
     fn with_new_children(
@@ -149,6 +153,7 @@ impl ExecutionPlan for HopExec {
             layers: self.layers.clone(),
             right_schema: self.right_schema.clone(),
             output_schema: self.output_schema.clone(),
+            props: self.props.clone(),
             right_proj: self.right_proj.clone(),
         }))
     }