Fix python

Author: kyotoYaho

python/Cargo.toml

@@ -37,6 +37,8 @@ default = ["mimalloc"]
 async-trait = "0.1"
 ballista = { path = "../ballista/client", version = "0.10.0" }
 datafusion = { version = "16.1.0", features = ["pyarrow"] }
+datafusion-common = "16.1.0"
+datafusion-expr = "16.1.0"
 futures = "0.3"
 mimalloc = { version = "*", optional = true, default-features = false }
 pyo3 = { version = "~0.17.1", features = ["extension-module", "abi3", "abi3-py37"] }

python/src/catalog.rs

@@ -27,6 +27,9 @@ use datafusion::{
     datasource::{TableProvider, TableType},
 };
 
+use crate::errors::DataFusionError;
+use crate::utils::wait_for_future;
+
 #[pyclass(name = "Catalog", module = "ballista", subclass)]
 pub(crate) struct PyCatalog {
     catalog: Arc<dyn CatalogProvider>,
@@ -88,13 +91,11 @@ impl PyDatabase {
         self.database.table_names().into_iter().collect()
     }
 
-    fn table(&self, name: &str) -> PyResult<PyTable> {
-        match self.database.table(name) {
-            Some(table) => Ok(PyTable::new(table)),
-            None => Err(PyKeyError::new_err(format!(
-                "Table with name {} doesn't exist.",
-                name
-            ))),
+    fn table(&self, name: &str, py: Python) -> PyResult<PyTable> {
+        if let Some(table) = wait_for_future(py, self.database.table(name)) {
+            Ok(PyTable::new(table))
+        } else {
+            Err(DataFusionError::Common(format!("Table with name {name} doesn't exist.")).into())
         }
     }
 

python/src/context.rs

@@ -29,7 +29,7 @@ use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::datasource::datasource::TableProvider;
 use datafusion::datasource::MemTable;
 use datafusion::execution::context::{SessionConfig, SessionContext};
-use datafusion::prelude::{CsvReadOptions, ParquetReadOptions};
+use datafusion::prelude::{CsvReadOptions, DataFrame, ParquetReadOptions};
 
 use crate::catalog::{PyCatalog, PyTable};
 use crate::dataframe::PyDataFrame;
@@ -75,7 +75,7 @@ impl PySessionContext {
         // TODO: config_options
     ) -> Self {
         let cfg = SessionConfig::new()
-            .create_default_catalog_and_schema(create_default_catalog_and_schema)
+            .with_create_default_catalog_and_schema(create_default_catalog_and_schema)
             .with_default_catalog_and_schema(default_catalog, default_schema)
             .with_information_schema(information_schema)
             .with_repartition_joins(repartition_joins)
@@ -103,6 +103,7 @@ impl PySessionContext {
     fn create_dataframe(
         &mut self,
         partitions: PyArrowType<Vec<Vec<RecordBatch>>>,
+        py: Python,
     ) -> PyResult<PyDataFrame> {
         let table = MemTable::try_new(partitions.0[0][0].schema(), partitions.0)
             .map_err(DataFusionError::from)?;
@@ -117,10 +118,9 @@ impl PySessionContext {
         self.ctx
             .register_table(&*name, Arc::new(table))
             .map_err(DataFusionError::from)?;
-        let table = self.ctx.table(&*name).map_err(DataFusionError::from)?;
+        let table = wait_for_future(py, self._table(&name)).map_err(DataFusionError::from)?;
 
-        let df = PyDataFrame::new(table);
-        Ok(df)
+        Ok(PyDataFrame::new(table))
     }
 
     fn register_table(&mut self, name: &str, table: &PyTable) -> PyResult<()> {
@@ -248,15 +248,23 @@ impl PySessionContext {
         self.ctx.tables().unwrap()
     }
 
-    fn table(&self, name: &str) -> PyResult<PyDataFrame> {
-        Ok(PyDataFrame::new(self.ctx.table(name)?))
+    fn table(&self, name: &str, py: Python) -> PyResult<PyDataFrame> {
+        let table = wait_for_future(py, self._table(name)).map_err(DataFusionError::from)?;
+
+        Ok(PyDataFrame::new(table))
     }
 
     fn empty_table(&self) -> PyResult<PyDataFrame> {
         Ok(PyDataFrame::new(self.ctx.read_empty()?))
     }
 }
 
+impl PySessionContext {
+    async fn _table(&self, name: &str) -> datafusion_common::Result<DataFrame> {
+        self.ctx.table(name).await
+    }
+}
+
 fn convert_table_partition_cols(
     table_partition_cols: Vec<(String, PyDataType)>,
 ) -> Vec<(String, DataType)> {

python/src/dataframe.rs

@@ -38,8 +38,8 @@ pub(crate) struct PyDataFrame {
 
 impl PyDataFrame {
     /// creates a new PyDataFrame
-    pub fn new(df: Arc<DataFrame>) -> Self {
-        Self { df }
+    pub fn new(df: DataFrame) -> Self {
+        Self { df: Arc::new(df) }
     }
 }
 
@@ -71,51 +71,51 @@ impl PyDataFrame {
 
     #[args(args = "*")]
     fn select_columns(&self, args: Vec<&str>) -> PyResult<Self> {
-        let df = self.df.select_columns(&args)?;
+        let df = self.df.as_ref().clone().select_columns(&args)?;
         Ok(Self::new(df))
     }
 
     #[args(args = "*")]
     fn select(&self, args: Vec<PyExpr>) -> PyResult<Self> {
         let expr = args.into_iter().map(|e| e.into()).collect();
-        let df = self.df.select(expr)?;
+        let df = self.df.as_ref().clone().select(expr)?;
         Ok(Self::new(df))
     }
 
     fn filter(&self, predicate: PyExpr) -> PyResult<Self> {
-        let df = self.df.filter(predicate.into())?;
+        let df = self.df.as_ref().clone().filter(predicate.into())?;
         Ok(Self::new(df))
     }
 
     fn with_column(&self, name: &str, expr: PyExpr) -> PyResult<Self> {
-        let df = self.df.with_column(name, expr.into())?;
+        let df = self.df.as_ref().clone().with_column(name, expr.into())?;
         Ok(Self::new(df))
     }
 
     fn aggregate(&self, group_by: Vec<PyExpr>, aggs: Vec<PyExpr>) -> PyResult<Self> {
         let group_by = group_by.into_iter().map(|e| e.into()).collect();
         let aggs = aggs.into_iter().map(|e| e.into()).collect();
-        let df = self.df.aggregate(group_by, aggs)?;
+        let df = self.df.as_ref().clone().aggregate(group_by, aggs)?;
         Ok(Self::new(df))
     }
 
     #[args(exprs = "*")]
     fn sort(&self, exprs: Vec<PyExpr>) -> PyResult<Self> {
         let exprs = exprs.into_iter().map(|e| e.into()).collect();
-        let df = self.df.sort(exprs)?;
+        let df = self.df.as_ref().clone().sort(exprs)?;
         Ok(Self::new(df))
     }
 
     fn limit(&self, count: usize) -> PyResult<Self> {
-        let df = self.df.limit(0, Some(count))?;
+        let df = self.df.as_ref().clone().limit(0, Some(count))?;
         Ok(Self::new(df))
     }
 
     /// Executes the plan, returning a list of `RecordBatch`es.
     /// Unless some order is specified in the plan, there is no
     /// guarantee of the order of the result.
     fn collect(&self, py: Python) -> PyResult<Vec<PyObject>> {
-        let batches = wait_for_future(py, self.df.collect())?;
+        let batches = wait_for_future(py, self.df.as_ref().clone().collect())?;
         // cannot use PyResult<Vec<RecordBatch>> return type due to
         // https://github.com/PyO3/pyo3/issues/1813
         batches.into_iter().map(|rb| rb.to_pyarrow(py)).collect()
@@ -124,7 +124,7 @@ impl PyDataFrame {
     /// Print the result, 20 lines by default
     #[args(num = "20")]
     fn show(&self, py: Python, num: usize) -> PyResult<()> {
-        let df = self.df.limit(0, Some(num))?;
+        let df = self.df.as_ref().clone().limit(0, Some(num))?;
         let batches = wait_for_future(py, df.collect())?;
         pretty::print_batches(&batches)
             .map_err(|err| PyArrowException::new_err(err.to_string()))
@@ -153,9 +153,13 @@ impl PyDataFrame {
             }
         };
 
-        let df = self
-            .df
-            .join(right.df, join_type, &join_keys.0, &join_keys.1, None)?;
+        let df = self.df.as_ref().clone().join(
+            right.df.as_ref().clone(),
+            join_type,
+            &join_keys.0,
+            &join_keys.1,
+            None,
+        )?;
         Ok(Self::new(df))
     }
 

python/src/dataset_exec.rs

@@ -162,10 +162,6 @@ impl ExecutionPlan for DatasetExec {
         })
     }
 
-    fn relies_on_input_order(&self) -> bool {
-        false
-    }
-
     fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
         None
     }