feat: update arrow to 51, datafusion to 37 (#2240)

Author: westonpace

Cargo.toml

@@ -57,17 +57,17 @@ lance-test-macros = { version = "=0.10.15", path = "./rust/lance-test-macros" }
 lance-testing = { version = "=0.10.15", path = "./rust/lance-testing" }
 approx = "0.5.1"
 # Note that this one does not include pyarrow
-arrow = { version = "50.0.0", optional = false, features = ["prettyprint"] }
-arrow-arith = "50.0"
-arrow-array = "50.0"
-arrow-buffer = "50.0"
-arrow-cast = "50.0"
-arrow-data = "50.0"
-arrow-ipc = { version = "50.0", features = ["zstd"] }
-arrow-ord = "50.0"
-arrow-row = "50.0"
-arrow-schema = "50.0"
-arrow-select = "50.0"
+arrow = { version = "51.0.0", optional = false, features = ["prettyprint"] }
+arrow-arith = "51.0"
+arrow-array = "51.0"
+arrow-buffer = "51.0"
+arrow-cast = "51.0"
+arrow-data = "51.0"
+arrow-ipc = { version = "51.0", features = ["zstd"] }
+arrow-ord = "51.0"
+arrow-row = "51.0"
+arrow-schema = "51.0"
+arrow-select = "51.0"
 async-recursion = "1.0"
 async-trait = "0.1"
 aws-config = "0.56"
@@ -85,14 +85,18 @@ chrono = { version = "0.4.25", default-features = false, features = [
     "now",
 ] }
 criterion = { version = "0.5", features = ["async", "async_tokio"] }
-datafusion = { version = "36.0.0", default-features = false, features = [
+datafusion = { version = "37.1", default-features = false, features = [
+    "array_expressions",
+    "regex_expressions",
+] }
+datafusion-common = "37.1"
+datafusion-functions = { version = "37.1", features = ["regex_expressions"] }
+datafusion-sql = "37.1"
+datafusion-expr = "37.1"
+datafusion-execution = "37.1"
+datafusion-physical-expr = { version = "37.1", features = [
     "regex_expressions",
 ] }
-datafusion-common = "36.0"
-datafusion-sql = "36.0"
-datafusion-expr = "36.0"
-datafusion-execution = "36.0"
-datafusion-physical-expr = "36.0"
 either = "1.0"
 futures = "0.3"
 http = "0.2.9"

python/Cargo.toml

@@ -12,10 +12,10 @@ name = "lance"
 crate-type = ["cdylib"]
 
 [dependencies]
-arrow = { version = "50.0.0", features = ["pyarrow"] }
-arrow-array = "50.0"
-arrow-data = "50.0"
-arrow-schema = "50.0"
+arrow = { version = "51.0.0", features = ["pyarrow"] }
+arrow-array = "51.0"
+arrow-data = "51.0"
+arrow-schema = "51.0"
 object_store = "0.9.0"
 async-trait = "0.1"
 chrono = "0.4.31"

rust/lance-datafusion/Cargo.toml

@@ -17,8 +17,9 @@ arrow-ord.workspace = true
 async-trait.workspace = true
 datafusion.workspace = true
 datafusion-common.workspace = true
+datafusion-functions.workspace = true
 datafusion-physical-expr.workspace = true
-datafusion-substrait = { version = "36.0", optional = true }
+datafusion-substrait = { version = "37.1", optional = true }
 futures.workspace = true
 lance-arrow.workspace = true
 lance-core = { workspace = true, features = ["datafusion"] }

rust/lance-datafusion/src/exec.rs

@@ -17,12 +17,12 @@ use datafusion::{
         TaskContext,
     },
     physical_plan::{
-        streaming::PartitionStream, DisplayAs, DisplayFormatType, ExecutionPlan,
+        streaming::PartitionStream, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
         SendableRecordBatchStream,
     },
 };
-use datafusion_common::DataFusionError;
-use datafusion_physical_expr::Partitioning;
+use datafusion_common::{DataFusionError, Statistics};
+use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
 
 use lance_arrow::SchemaExt;
 use lance_core::Result;
@@ -32,11 +32,15 @@ use log::{info, warn};
 ///
 /// It can only be used once, and will return the stream.  After that the node
 /// is exhuasted.
+///
+/// Note: the stream should be finite, otherwise we will report datafusion properties
+/// incorrectly.
 pub struct OneShotExec {
     stream: Mutex<Option<SendableRecordBatchStream>>,
     // We save off a copy of the schema to speed up formatting and so ExecutionPlan::schema & display_as
     // can still function after exhuasted
     schema: Arc<ArrowSchema>,
+    properties: PlanProperties,
 }
 
 impl OneShotExec {
@@ -45,7 +49,12 @@ impl OneShotExec {
         let schema = stream.schema().clone();
         Self {
             stream: Mutex::new(Some(stream)),
-            schema,
+            schema: schema.clone(),
+            properties: PlanProperties::new(
+                EquivalenceProperties::new(schema),
+                Partitioning::RoundRobinBatch(1),
+                datafusion::physical_plan::ExecutionMode::Bounded,
+            ),
         }
     }
 }
@@ -96,14 +105,6 @@ impl ExecutionPlan for OneShotExec {
         self.schema.clone()
     }
 
-    fn output_partitioning(&self) -> datafusion_physical_expr::Partitioning {
-        Partitioning::RoundRobinBatch(1)
-    }
-
-    fn output_ordering(&self) -> Option<&[datafusion_physical_expr::PhysicalSortExpr]> {
-        None
-    }
-
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![]
     }
@@ -135,7 +136,11 @@ impl ExecutionPlan for OneShotExec {
     }
 
     fn statistics(&self) -> datafusion_common::Result<datafusion_common::Statistics> {
-        todo!()
+        Ok(Statistics::new_unknown(&self.schema))
+    }
+
+    fn properties(&self) -> &datafusion::physical_plan::PlanProperties {
+        &self.properties
     }
 }
 
@@ -194,7 +199,7 @@ pub fn execute_plan(
     let session_state = SessionState::new_with_config_rt(session_config, runtime_env);
     // NOTE: we are only executing the first partition here. Therefore, if
     // the plan has more than one partition, we will be missing data.
-    assert_eq!(plan.output_partitioning().partition_count(), 1);
+    assert_eq!(plan.properties().partitioning.partition_count(), 1);
     Ok(plan.execute(0, session_state.task_ctx())?)
 }
 

rust/lance-datafusion/src/expr.rs

@@ -508,7 +508,7 @@ pub async fn parse_substrait(expr: &[u8], input_schema: Arc<Schema>) -> Result<E
                 match relation {
                     TableReference::Bare { table } => {
                         if table == "dummy" {
-                            Ok(Transformed::Yes(Expr::Column(Column {
+                            Ok(Transformed::yes(Expr::Column(Column {
                                 relation: None,
                                 name: column.name,
                             })))
@@ -524,12 +524,12 @@ pub async fn parse_substrait(expr: &[u8], input_schema: Arc<Schema>) -> Result<E
                             _ => Err(DataFusionError::Substrait("Unexpected partially or fully qualified table reference encountered when parsing filter".into()))
                 }
             } else {
-                Ok(Transformed::No(Expr::Column(column)))
+                Ok(Transformed::no(Expr::Column(column)))
             }
         }
-        _ => Ok(Transformed::No(node)),
+        _ => Ok(Transformed::no(node)),
     })?;
-    Ok(expr)
+    Ok(expr.data)
 }
 
 #[cfg(test)]

rust/lance-index/src/scalar/btree.rs

@@ -494,6 +494,8 @@ impl Ord for OrderableScalarValue {
             (Dictionary(_k1, _v1), Dictionary(_k2, _v2)) => todo!(),
             (Dictionary(_, v1), Null) => Self(*v1.clone()).cmp(&Self(ScalarValue::Null)),
             (Dictionary(_, _), _) => panic!("Attempt to compare Dictionary with non-Dictionary"),
+            // What would a btree of unions even look like?  May not be possible.
+            (Union(_, _, _), _) => todo!("Support for union scalars"),
             (Null, Null) => Ordering::Equal,
             (Null, _) => todo!(),
         }

rust/lance-index/src/scalar/expression.rs

@@ -549,6 +549,18 @@ mod tests {
         fn options(&self) -> &ConfigOptions {
             todo!()
         }
+
+        fn udfs_names(&self) -> Vec<String> {
+            todo!()
+        }
+
+        fn udafs_names(&self) -> Vec<String> {
+            todo!()
+        }
+
+        fn udwfs_names(&self) -> Vec<String> {
+            todo!()
+        }
     }
 
     fn check(

rust/lance/Cargo.toml

@@ -57,6 +57,7 @@ arrow.workspace = true
 num_cpus.workspace = true
 # TODO: use datafusion sub-modules to reduce build size?
 datafusion.workspace = true
+datafusion-functions.workspace = true
 datafusion-physical-expr.workspace = true
 lapack = { version = "0.19.0", optional = true }
 lru_time_cache = "0.11"