Fix python

apache · Dec 8, 2022 · 932ed45 · 932ed45
1 parent 8e52f1b
commit 932ed45
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 8 deletions.
diff --git a/python/src/context.rs b/python/src/context.rs
@@ -23,9 +23,7 @@ use uuid::Uuid;
 use pyo3::exceptions::{PyKeyError, PyValueError};
 use pyo3::prelude::*;
 
-use arrow::datatypes::DataType;
-
-use datafusion::arrow::datatypes::Schema;
+use datafusion::arrow::datatypes::{DataType, Schema};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::datasource::datasource::TableProvider;
@@ -161,13 +159,13 @@ impl PySessionContext {
         &mut self,
         name: &str,
         path: &str,
-        table_partition_cols: Vec<(String, DataType)>,
+        table_partition_cols: Vec<(String, PyDataType)>,
         parquet_pruning: bool,
         file_extension: &str,
         py: Python,
     ) -> PyResult<()> {
         let mut options = ParquetReadOptions::default()
-            .table_partition_cols(table_partition_cols)
+            .table_partition_cols(convert_table_partition_cols(table_partition_cols))
             .parquet_pruning(parquet_pruning);
         options.file_extension = file_extension;
         let result = self.ctx.register_parquet(name, path, options);
@@ -257,3 +255,12 @@ impl PySessionContext {
         Ok(PyDataFrame::new(self.ctx.read_empty()?))
     }
 }
+
+fn convert_table_partition_cols(
+    table_partition_cols: Vec<(String, PyDataType)>,
+) -> Vec<(String, DataType)> {
+    table_partition_cols
+        .iter()
+        .map(|(name, t)| (name.clone(), t.data_type.clone()))
+        .collect()
+}
diff --git a/python/src/dataset.rs b/python/src/dataset.rs
@@ -111,7 +111,7 @@ impl TableProvider for Dataset {
                 DatasetExec::new(
                     py,
                     self.dataset.as_ref(py),
-                    projection.clone(),
+                    projection.cloned(),
                     filters,
                 )
                 .map_err(|err| DataFusionError::External(Box::new(err)))?,

diff --git a/python/src/udaf.rs b/python/src/udaf.rs
@@ -97,8 +97,7 @@ impl Accumulator for RustAccumulator {
     }
 
     fn size(&self) -> usize {
-        Python::with_gil(|py| self.accum.as_ref(py).call_method0("size")?.extract())
-            .map_err(|e| DataFusionError::Execution(format!("{}", e)))
+        std::mem::size_of_val(self)
     }
 }