WIP update to main on datafusion

Author: timsaucer

Cargo.lock

@@ -39,10 +39,10 @@ pyo3 = { version = "0.25", features = ["extension-module", "abi3", "abi3-py39"]
 pyo3-async-runtimes = { version = "0.25", features = ["tokio-runtime"]}
 pyo3-log = "0.12.4"
 arrow = { version = "56", features = ["pyarrow"] }
-datafusion = { version = "50", features = ["avro", "unicode_expressions"] }
-datafusion-substrait = { version = "50", optional = true }
-datafusion-proto = { version = "50" }
-datafusion-ffi = { version = "50" }
+datafusion = { version = "50.3.0", features = ["avro", "unicode_expressions"] }
+datafusion-substrait = { version = "50.3.0", optional = true }
+datafusion-proto = { version = "50.3.0" }
+datafusion-ffi = { version = "50.3.0" }
 prost = "0.13.1" # keep in line with `datafusion-substrait`
 uuid = { version = "1.18", features = ["v4"] }
 mimalloc = { version = "0.1", optional = true, default-features = false, features = ["local_dynamic_tls"] }
@@ -64,3 +64,9 @@ crate-type = ["cdylib", "rlib"]
 [profile.release]
 lto = true
 codegen-units = 1
+
+[patch.crates-io]
+datafusion = { git = "https://github.com/apache/datafusion", branch = "main" }
+datafusion-substrait = { git = "https://github.com/apache/datafusion", branch = "main" }
+datafusion-proto = { git = "https://github.com/apache/datafusion", branch = "main" }
+datafusion-ffi = { git = "https://github.com/apache/datafusion", branch = "main" }

Cargo.toml

@@ -18,7 +18,7 @@
 use datafusion::arrow::array::Array;
 use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
 use datafusion::common::ScalarValue;
-use datafusion::logical_expr::sqlparser::ast::NullTreatment as DFNullTreatment;
+use datafusion::logical_expr::expr::NullTreatment as DFNullTreatment;
 use pyo3::exceptions::PyNotImplementedError;
 use pyo3::{exceptions::PyValueError, prelude::*};
 
@@ -261,6 +261,12 @@ impl DataTypeMap {
             ScalarValue::Float16(_) => Ok(DataType::Float16),
             ScalarValue::Float32(_) => Ok(DataType::Float32),
             ScalarValue::Float64(_) => Ok(DataType::Float64),
+            ScalarValue::Decimal32(_, precision, scale) => {
+                Ok(DataType::Decimal32(*precision, *scale))
+            }
+            ScalarValue::Decimal64(_, precision, scale) => {
+                Ok(DataType::Decimal64(*precision, *scale))
+            }
             ScalarValue::Decimal128(_, precision, scale) => {
                 Ok(DataType::Decimal128(*precision, *scale))
             }

src/common/data_type.rs

@@ -412,7 +412,7 @@ impl PyDataFrame {
 
     /// Returns the schema from the logical plan
     fn schema(&self) -> PyArrowType<Schema> {
-        PyArrowType(self.df.schema().into())
+        PyArrowType(self.df.schema().as_arrow().to_owned())
     }
 
     /// Convert this DataFrame into a Table Provider that can be used in register_table
@@ -925,7 +925,7 @@ impl PyDataFrame {
         requested_schema: Option<Bound<'py, PyCapsule>>,
     ) -> PyDataFusionResult<Bound<'py, PyCapsule>> {
         let mut batches = wait_for_future(py, self.df.as_ref().clone().collect())??;
-        let mut schema: Schema = self.df.schema().to_owned().into();
+        let mut schema: Schema = self.df.schema().as_arrow().clone();
 
         if let Some(schema_capsule) = requested_schema {
             validate_pycapsule(&schema_capsule, "arrow_schema")?;

src/dataframe.rs

@@ -66,14 +66,15 @@ impl Display for PyCreateExternalTable {
 impl PyCreateExternalTable {
     #[allow(clippy::too_many_arguments)]
     #[new]
-    #[pyo3(signature = (schema, name, location, file_type, table_partition_cols, if_not_exists, temporary, order_exprs, unbounded, options, constraints, column_defaults, definition=None))]
+    #[pyo3(signature = (schema, name, location, file_type, table_partition_cols, if_not_exists, or_replace, temporary, order_exprs, unbounded, options, constraints, column_defaults, definition=None))]
     pub fn new(
         schema: PyDFSchema,
         name: String,
         location: String,
         file_type: String,
         table_partition_cols: Vec<String>,
         if_not_exists: bool,
+        or_replace: bool,
         temporary: bool,
         order_exprs: Vec<Vec<PySortExpr>>,
         unbounded: bool,
@@ -89,6 +90,7 @@ impl PyCreateExternalTable {
             file_type,
             table_partition_cols,
             if_not_exists,
+            or_replace,
             temporary,
             definition,
             order_exprs: order_exprs

src/expr/create_external_table.rs

@@ -38,10 +38,11 @@ impl PyPlaceholder {
         self.placeholder.id.clone()
     }
 
+    // TODO now that the inner struct holds field we should expose it instead
     fn data_type(&self) -> Option<PyDataType> {
         self.placeholder
-            .data_type
+            .field
             .as_ref()
-            .map(|e| e.clone().into())
+            .map(|f| f.data_type().clone().into())
     }
 }

src/expr/placeholder.rs

@@ -15,11 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::datatypes::{DataType, Field};
 use datafusion::logical_expr::{
     Deallocate, Execute, Prepare, SetVariable, TransactionAccessMode, TransactionConclusion,
     TransactionEnd, TransactionIsolationLevel, TransactionStart,
 };
 use pyo3::{prelude::*, IntoPyObjectExt};
+use std::sync::Arc;
 
 use crate::{common::data_type::PyDataType, sql::logical::PyLogicalPlan};
 
@@ -334,17 +336,22 @@ impl LogicalNode for PyPrepare {
 
 #[pymethods]
 impl PyPrepare {
+    // TODO Now that we have field instead of datatype in the inner struct
+    // we should update the signatures below to match
     #[new]
     pub fn new(name: String, data_types: Vec<PyDataType>, input: PyLogicalPlan) -> Self {
         let input = input.plan().clone();
-        let data_types = data_types
+        let fields = data_types
             .into_iter()
-            .map(|data_type| data_type.into())
+            .map(|data_type| DataType::from(data_type))
+            .enumerate()
+            .map(|(idx, data_type)| Field::new(format!("field_{idx}"), data_type, true))
+            .map(Arc::new)
             .collect();
         PyPrepare {
             prepare: Prepare {
                 name,
-                data_types,
+                fields,
                 input,
             },
         }
@@ -356,10 +363,9 @@ impl PyPrepare {
 
     pub fn data_types(&self) -> Vec<PyDataType> {
         self.prepare
-            .data_types
-            .clone()
-            .into_iter()
-            .map(|t| t.into())
+            .fields
+            .iter()
+            .map(|f| f.data_type().to_owned().into())
             .collect()
     }
 

src/expr/statement.rs

@@ -41,7 +41,7 @@ use datafusion::functions;
 use datafusion::functions_aggregate;
 use datafusion::functions_window;
 use datafusion::logical_expr::expr::Alias;
-use datafusion::logical_expr::sqlparser::ast::NullTreatment as DFNullTreatment;
+use datafusion::logical_expr::expr::NullTreatment as DFNullTreatment;
 use datafusion::logical_expr::{expr::WindowFunction, lit, Expr, WindowFunctionDefinition};
 
 fn add_builder_fns_to_aggregate(

src/functions.rs

@@ -83,7 +83,7 @@ impl PyExecutionPlan {
             })?;
 
         let codec = DefaultPhysicalExtensionCodec {};
-        let plan = proto_plan.try_into_physical_plan(&ctx.ctx, &ctx.ctx.runtime_env(), &codec)?;
+        let plan = proto_plan.try_into_physical_plan(&ctx.ctx.task_ctx(), &codec)?;
         Ok(Self::new(plan))
     }
 

src/physical_plan.rs

@@ -206,7 +206,7 @@ impl PyLogicalPlan {
             })?;
 
         let codec = DefaultLogicalExtensionCodec {};
-        let plan = proto_plan.try_into_logical_plan(&ctx.ctx, &codec)?;
+        let plan = proto_plan.try_into_logical_plan(&ctx.ctx.task_ctx(), &codec)?;
         Ok(Self::new(plan))
     }
 }