update

Author: chenkovsky

python/datafusion/common.py

@@ -33,6 +33,9 @@
 SqlTable = common_internal.SqlTable
 SqlType = common_internal.SqlType
 SqlView = common_internal.SqlView
+TableType = common_internal.TableType
+TableSource = common_internal.TableSource
+Constraints = common_internal.Constraints
 
 __all__ = [
     "DFSchema",
@@ -47,6 +50,9 @@
     "SqlTable",
     "SqlType",
     "SqlView",
+    "TableType",
+    "TableSource",
+    "Constraints",
 ]
 
 

src/common.rs

@@ -36,5 +36,8 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<schema::SqlView>()?;
     m.add_class::<schema::SqlStatistics>()?;
     m.add_class::<function::SqlFunction>()?;
+    m.add_class::<schema::PyTableType>()?;
+    m.add_class::<schema::PyTableSource>()?;
+    m.add_class::<schema::PyConstraints>()?;
     Ok(())
 }

src/common/schema.rs

@@ -15,14 +15,22 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::fmt::{self, Display, Formatter};
+use std::sync::Arc;
 use std::{any::Any, borrow::Cow};
 
+use arrow::datatypes::Schema;
+use arrow::pyarrow::PyArrowType;
 use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::common::Constraints;
+use datafusion::datasource::TableType;
 use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableSource};
 use pyo3::prelude::*;
 
 use datafusion::logical_expr::utils::split_conjunction;
 
+use crate::sql::logical::PyLogicalPlan;
+
 use super::{data_type::DataTypeMap, function::SqlFunction};
 
 #[pyclass(name = "SqlSchema", module = "datafusion.common", subclass)]
@@ -218,3 +226,84 @@ impl SqlStatistics {
         self.row_count
     }
 }
+
+#[pyclass(name = "Constraints", module = "datafusion.expr", subclass)]
+#[derive(Clone)]
+pub struct PyConstraints {
+    pub constraints: Constraints,
+}
+
+impl From<PyConstraints> for Constraints {
+    fn from(constraints: PyConstraints) -> Self {
+        constraints.constraints
+    }
+}
+
+impl From<Constraints> for PyConstraints {
+    fn from(constraints: Constraints) -> Self {
+        PyConstraints { constraints }
+    }
+}
+
+impl Display for PyConstraints {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        write!(f, "Constraints: {:?}", self.constraints)
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[pyclass(eq, eq_int, name = "TableType", module = "datafusion.common")]
+pub enum PyTableType {
+    Base,
+    View,
+    Temporary,
+}
+
+impl From<PyTableType> for datafusion::logical_expr::TableType {
+    fn from(table_type: PyTableType) -> Self {
+        match table_type {
+            PyTableType::Base => datafusion::logical_expr::TableType::Base,
+            PyTableType::View => datafusion::logical_expr::TableType::View,
+            PyTableType::Temporary => datafusion::logical_expr::TableType::Temporary,
+        }
+    }
+}
+
+impl From<TableType> for PyTableType {
+    fn from(table_type: TableType) -> Self {
+        match table_type {
+            datafusion::logical_expr::TableType::Base => PyTableType::Base,
+            datafusion::logical_expr::TableType::View => PyTableType::View,
+            datafusion::logical_expr::TableType::Temporary => PyTableType::Temporary,
+        }
+    }
+}
+
+#[pyclass(name = "TableSource", module = "datafusion.common", subclass)]
+#[derive(Clone)]
+pub struct PyTableSource {
+    pub table_source: Arc<dyn TableSource>,
+}
+
+#[pymethods]
+impl PyTableSource {
+    pub fn schema(&self) -> PyArrowType<Schema> {
+        (*self.table_source.schema()).clone().into()
+    }
+
+    pub fn constraints(&self) -> Option<PyConstraints> {
+        self.table_source.constraints().map(|c| PyConstraints {
+            constraints: c.clone(),
+        })
+    }
+
+    pub fn table_type(&self) -> PyTableType {
+        self.table_source.table_type().into()
+    }
+
+    pub fn get_logical_plan(&self) -> Option<PyLogicalPlan> {
+        self.table_source
+            .get_logical_plan()
+            .map(|plan| PyLogicalPlan::new(plan.into_owned()))
+    }
+}

src/expr.rs

@@ -66,7 +66,6 @@ pub mod case;
 pub mod cast;
 pub mod column;
 pub mod conditional_expr;
-pub mod constraints;
 pub mod copy_to;
 pub mod create_catalog;
 pub mod create_catalog_schema;
@@ -810,7 +809,6 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<create_catalog::PyCreateCatalog>()?;
     m.add_class::<create_catalog_schema::PyCreateCatalogSchema>()?;
     m.add_class::<create_external_table::PyCreateExternalTable>()?;
-    m.add_class::<constraints::PyConstraints>()?;
     m.add_class::<create_function::PyCreateFunction>()?;
     m.add_class::<create_function::PyOperateFunctionArg>()?;
     m.add_class::<create_function::PyCreateFunctionBody>()?;

src/expr/constraints.rs

@@ -22,7 +22,7 @@ use std::{
 };
 
 use datafusion::{common::file_options::file_type::FileType, logical_expr::dml::CopyTo};
-use pyo3::prelude::*;
+use pyo3::{prelude::*, IntoPyObjectExt};
 
 use crate::sql::logical::PyLogicalPlan;
 
@@ -57,8 +57,8 @@ impl LogicalNode for PyCopyTo {
         vec![PyLogicalPlan::from((*self.copy.input).clone())]
     }
 
-    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
-        Ok(self.clone().into_py(py))
+    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        self.clone().into_bound_py_any(py)
     }
 }
 

src/expr/copy_to.rs

@@ -21,7 +21,7 @@ use std::{
 };
 
 use datafusion::logical_expr::CreateCatalog;
-use pyo3::prelude::*;
+use pyo3::{prelude::*, IntoPyObjectExt};
 
 use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
 
@@ -94,7 +94,7 @@ impl LogicalNode for PyCreateCatalog {
         vec![]
     }
 
-    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
-        Ok(self.clone().into_py(py))
+    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        self.clone().into_bound_py_any(py)
     }
 }

src/expr/create_catalog.rs

@@ -21,7 +21,7 @@ use std::{
 };
 
 use datafusion::logical_expr::CreateCatalogSchema;
-use pyo3::prelude::*;
+use pyo3::{prelude::*, IntoPyObjectExt};
 
 use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
 
@@ -94,7 +94,7 @@ impl LogicalNode for PyCreateCatalogSchema {
         vec![]
     }
 
-    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
-        Ok(self.clone().into_py(py))
+    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        self.clone().into_bound_py_any(py)
     }
 }

src/expr/create_catalog_schema.rs

@@ -15,19 +15,19 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::{expr::PyExpr, sql::logical::PyLogicalPlan};
+use crate::{common::schema::PyConstraints, expr::PyExpr, sql::logical::PyLogicalPlan};
 use std::{
     collections::HashMap,
     fmt::{self, Display, Formatter},
     sync::Arc,
 };
 
 use datafusion::logical_expr::CreateExternalTable;
-use pyo3::prelude::*;
+use pyo3::{prelude::*, IntoPyObjectExt};
 
 use crate::common::df_schema::PyDFSchema;
 
-use super::{constraints::PyConstraints, logical_node::LogicalNode, sort_expr::PySortExpr};
+use super::{logical_node::LogicalNode, sort_expr::PySortExpr};
 
 #[pyclass(name = "CreateExternalTable", module = "datafusion.expr", subclass)]
 #[derive(Clone)]
@@ -59,6 +59,7 @@ impl Display for PyCreateExternalTable {
 
 #[pymethods]
 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))]
     pub fn new(
@@ -176,7 +177,7 @@ impl LogicalNode for PyCreateExternalTable {
         vec![]
     }
 
-    fn to_variant(&self, py: Python<'_>) -> PyResult<PyObject> {
-        Ok(self.clone().into_py(py))
+    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        self.clone().into_bound_py_any(py)
     }
 }

src/expr/create_external_table.rs

@@ -23,7 +23,7 @@ use std::{
 use datafusion::logical_expr::{
     CreateFunction, CreateFunctionBody, OperateFunctionArg, Volatility,
 };
-use pyo3::prelude::*;
+use pyo3::{prelude::*, IntoPyObjectExt};
 
 use super::logical_node::LogicalNode;
 use super::PyExpr;
@@ -176,7 +176,7 @@ impl LogicalNode for PyCreateFunction {
         vec![]
     }
 
-    fn to_variant(&self, py: Python<'_>) -> PyResult<PyObject> {
-        Ok(self.clone().into_py(py))
+    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        self.clone().into_bound_py_any(py)
     }
 }