8000 update · chenkovsky/datafusion-python@ba4696f · GitHub
[go: up one dir, main page]
Skip to content

Commit ba4696f

Browse files
chenkovskychenkovsky
committed
update
1 parent d4e3420 commit ba4696f

19 files changed

+161
-103
lines changed

‎python/datafusion/common.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@
3333
SqlTable = common_internal.SqlTable
3434
SqlType = common_internal.SqlType
3535
SqlView = common_internal.SqlView
36+
TableType = common_internal.TableType
37+
TableSource = common_internal.TableSource
38+
Constraints = common_internal.Constraints
3639

3740
__all__ = [
3841
"DFSchema",
@@ -47,6 +50,9 @@
4750
"SqlTable" A3E2 ;,
4851
"SqlType",
4952
"SqlView",
53+
"TableType",
54+
"TableSource",
55+
"Constraints",
5056
]
5157

5258

‎src/common.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,5 +36,8 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
3636
m.add_class::<schema::SqlView>()?;
3737
m.add_class::<schema::SqlStatistics>()?;
3838
m.add_class::<function::SqlFunction>()?;
39+
m.add_class::<schema::PyTableType>()?;
40+
m.add_class::<schema::PyTableSource>()?;
41+
m.add_class::<schema::PyConstraints>()?;
3942
Ok(())
4043
}

‎src/common/schema.rs

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,22 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use std::fmt::{self, Display, Formatter};
19+
use std::sync::Arc;
1820
use std::{any::Any, borrow::Cow};
1921

22+
use arrow::datatypes::Schema;
23+
use arrow::pyarrow::PyArrowType;
2024
use datafusion::arrow::datatypes::SchemaRef;
25+
use datafusion::common::Constraints;
26+
use datafusion::datasource::TableType;
2127
use datafusion::logical_expr::{Expr, TableProviderFilterPushDown, TableSource};
2228
use pyo3::prelude::*;
2329

2430
use datafusion::logical_expr::utils::split_conjunction;
2531

32+
use crate::sql::logical::PyLogicalPlan;
33+
2634
use super::{data_type::DataTypeMap, function::SqlFunction};
2735

2836
#[pyclass(name = "SqlSchema", module = "datafusion.common", subclass)]
@@ -218,3 +226,84 @@ impl SqlStatistics {
218226
self.row_count
219227
}
220228
}
229+
230+
#[pyclass(name = "Constraints", module = "datafusion.expr", subclass)]
231+
#[derive(Clone)]
232+
pub struct PyConstraints {
233+
pub constraints: Constraints,
234+
}
235+
236+
impl From<PyConstraints> for Constraints {
237+
fn from(constraints: PyConstraints) -> Self {
238+
constraints.constraints
239+
}
240+
}
241+
242+
impl From<Constraints> for PyConstraints {
243+
fn from(constraints: Constraints) -> Self {
244+
PyConstraints { constraints }
245+
}
246+
}
247+
248+
impl Display for PyConstraints {
249+
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
250+
write!(f, "Constraints: {:?}", self.constraints)
251+
}
252+
}
253+
254+
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
255+
#[pyclass(eq, eq_int, name = "TableType", module = "datafusion.common")]
256+
pub enum PyTableType {
257+
Base,
258+
View,
259+
Temporary,
260+
}
261+
262+
impl From<PyTableType> for datafusion::logical_expr::TableType {
263+
fn from(table_type: PyTableType) -> Self {
264+
match table_type {
265+
PyTableType::Base => datafusion::logical_expr::TableType::Base,
266+
PyTableType::View => datafusion::logical_expr::TableType::View,
267+
PyTableType::Temporary => datafusion::logical_expr::TableType::Temporary,
268+
}
269+
}
270+
}
271+
272+
impl From<TableType> for PyTableType {
273+
fn from(table_type: TableType) -> Self {
274+
match table_type {
275+
datafusion::logical_expr::TableType::Base => PyTableType::Base,
276+
datafusion::logical_expr::TableType::View => PyTableType::View,
277+
datafusion::logical_expr::TableType::Temporary => PyTableType::Temporary,
278+
}
279+
}
280+
}
281+
282+
#[pyclass(name = "TableSource", module = "datafusion.common", subclass)]
283+
#[derive(Clone)]
284+
pub struct PyTableSource {
285+
pub table_source: Arc<dyn TableSource>,
286+
}
287+
288+
#[pymethods]
289+
impl PyTableSource {
290+
pub fn schema(&self) -> PyArrowType<Schema> {
291+
(*self.table_source.schema()).clone().into()
292+
}
293+
294+
pub fn constraints(&self) -> Option<PyConstraints> {
295+
self.table_source.constraints().map(|c| PyConstraints {
296+
constraints: c.clone(),
297+
})
298+
}
299+
300+
pub fn table_type(&self) -> PyTableType {
301+
self.table_source.table_type().into()
302+
}
303+
304+
pub fn get_logical_plan(&self) -> Option<PyLogicalPlan> {
305+
self.table_source
306+
.get_logical_plan()
307+
.map(|plan| PyLogicalPlan::new(plan.into_owned()))
308+
}
309+
}

‎src/expr.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,6 @@ pub mod case;
6666
pub mod cast;
6767
pub mod column;
6868
pub mod conditional_expr;
69-
pub mod constraints;
7069
pub mod copy_to;
7170
pub mod create_catalog;
7271
pub mod create_catalog_schema;
@@ -810,7 +809,6 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
810809
m.add_class::<create_catalog::PyCreateCatalog>()?;
811810
m.add_class::<create_catalog_schema::PyCreateCatalogSchema>()?;
812811
m.add_class::<create_external_table::PyCreateExternalTable>()?;
813-
m.add_class::<constraints::PyConstraints>()?;
814812
m.add_class::<create_function::PyCreateFunction>()?;
815813
m.add_class::<create_function::PyOperateFunctionArg>()?;
816814
m.add_class::<create_function::PyCreateFunctionBody>()?;

‎src/expr/constraints.rs

Lines changed: 0 additions & 45 deletions
This file was deleted.

‎src/expr/copy_to.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ use std::{
2222
};
2323

2424
use datafusion::{common::file_options::file_type::FileType, logical_expr::dml::CopyTo};
25-
use pyo3::prelude::*;
25+
use pyo3::{prelude::*, IntoPyObjectExt};
2626

2727
use crate::sql::logical::PyLogicalPlan;
2828

@@ -57,8 +57,8 @@ impl LogicalNode for PyCopyTo {
5757
vec![PyLogicalPlan::from((*self.copy.input).clone())]
5858
}
5959

60-
fn to_variant(&self, py: Python) -> PyResult<PyObject> {
61-
Ok(self.clone().into_py(py))
60+
fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
61+
self.clone().into_bound_py_any(py)
6262
}
6363
}
6464

‎src/expr/create_catalog.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use std::{
2121
};
2222

2323
use datafusion::logical_expr::CreateCatalog;
24-
use pyo3::prelude::*;
24+
use pyo3::{prelude::*, IntoPyObjectExt};
2525

2626
use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
2727

@@ -94,7 +94,7 @@ impl LogicalNode for PyCreateCatalog {
9494
vec![]
9595
}
9696

97-
fn to_variant(&self, py: Python) -> PyResult<PyObject> {
98-
Ok(self.clone().into_py(py))
97+
fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
98+
self.clone().into_bound_py_any(py)
9999
}
100100
}

‎src/expr/create_catalog_schema.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use std::{
2121
};
2222

2323
use datafusion::logical_expr::CreateCatalogSchema;
24-
use pyo3::prelude::*;
24+
use pyo3::{prelude::*, IntoPyObjectExt};
2525

2626
use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
2727

@@ -94,7 +94,7 @@ impl LogicalNode for PyCreateCatalogSchema {
9494
vec![]
9595
}
9696

97-
fn to_variant(&self, py: Python) -> PyResult<PyObject> {
98-
Ok(self.clone().into_py(py))
97+
fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
98+
self.clone().into_bound_py_any(py)
9999
}
100100
}

‎src/expr/create_external_table.rs

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,19 +15,19 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use crate::{expr::PyExpr, sql::logical::PyLogicalPlan};
18+
use crate::{common::schema::PyConstraints, expr::PyExpr, sql::logical::PyLogicalPlan};
1919
use std::{
2020
collections::HashMap,
2121
fmt::{self, Display, Formatter},
2222
sync::Arc,
2323
};
2424

2525
use datafusion::logical_expr::CreateExternalTable;
26-
use pyo3::prelude::*;
26+
use pyo3::{prelude::*, IntoPyObjectExt};
2727

2828
use crate::common::df_schema::PyDFSchema;
2929

30-
use super::{constraints::PyConstraints, logical_node::LogicalNode, sort_expr::PySortExpr};
30+
use super::{logical_node::LogicalNode, sort_expr::PySortExpr};
3131

3232
#[pyclass(name = "CreateExternalTable", module = "datafusion.expr", subclass)]
3333
#[derive(Clone)]
@@ -59,6 +59,7 @@ impl Display for PyCreateExternalTable {
5959

6060
#[pymethods]
6161
impl PyCreateExternalTable {
62+
#[allow(clippy::too_many_arguments)]
6263
#[new]
6364
#[pyo3(signature = (schema, name, location, file_type, table_partition_cols, if_not_exists, temporary, order_exprs, unbounded, options, constraints, column_defaults, definition=None))]
6465
pub fn new(
@@ -176,7 +177,7 @@ impl LogicalNode for PyCreateExternalTable {
176177
vec![]
177178
}
178179

179-
fn to_variant(&self, py: Python<'_>) -> PyResult<PyObject> {
180-
Ok(self.clone().into_py(py))
180+
fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
181+
self.clone().into_bound_py_any(py)
181182
}
182183
}

‎src/expr/create_function.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ use std::{
2323
use datafusion::logical_expr::{
2424
CreateFunction, CreateFunctionBody, OperateFunctionArg, Volatility,
2525
};
26-
use pyo3::prelude::*;
26+
use pyo3::{prelude::*, IntoPyObjectExt};
2727

2828
use super::logical_node::LogicalNode;
2929
use super::PyExpr;
@@ -176,7 +176,7 @@ impl LogicalNode for PyCreateFunction {
176176
vec![]
177177
}
178178

179-
fn to_variant(&self, py: Python<'_>) -> PyResult<PyObject> {
180-
Ok(self.clone().into_py(py))
179+
fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
180+
self.clone().into_bound_py_any(py)
181181
}
182182
}

‎src/expr/create_index.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use std::{
2121
};
2222

2323
use datafusion::logical_expr::CreateIndex;
24-
use pyo3::prelude::*;
24+
use pyo3::{prelude::*, IntoPyObjectExt};
2525

2626
use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
2727

@@ -123,7 +123,7 @@ impl LogicalNode for PyCreateIndex {
123123
vec![]
124124
}
125125

126-
fn to_variant(&self, py: Python) -> PyResult<PyObject> {
127-
Ok(self.clone().into_py(py))
126+
fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
127+
self.clone().into_bound_py_any(py)
128128
}
129129
}

‎src/expr/describe_table.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ use std::{
2222

2323
use arrow::{datatypes::Schema, pyarrow::PyArrowType};
2424
use datafusion::logical_expr::DescribeTable;
25-
use pyo3::prelude::*;
25+
use pyo3::{prelude::*, IntoPyObjectExt};
2626

2727
use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
2828

@@ -86,7 +86,7 @@ impl LogicalNode for PyDescribeTable {
8686
vec![]
8787
}
8888

89-
fn to_variant(&self, py: Python<'_>) -> PyResult<PyObject> {
90-
Ok(self.clone().into_py(py))
89+
fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
90+
self.clone().into_bound_py_any(py)
9191
}
9292
}

‎src/expr/dml.rs

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,9 @@
1717

1818
use datafusion::logical_expr::dml::InsertOp;
1919
use datafusion::logical_expr::{DmlStatement, WriteOp};
20-
use pyo3::prelude::*;
20+
use pyo3::{prelude::*, IntoPyObjectExt};
2121

22+
use crate::common::schema::PyTableSource;
2223
use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
2324

2425
use super::logical_node::LogicalNode;
@@ -46,8 +47,8 @@ impl LogicalNode for PyDmlStatement {
4647
vec![PyLogicalPlan::from((*self.dml.input).clone())]
4748
}
4849

49-
fn to_variant(&self, py: Python<'_>) -> PyResult<PyObject> {
50-
Ok(self.clone().into_py(py))
50+
fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
51+
self.clone().into_bound_py_any(py)
5152
}
5253
}
5354

@@ -57,8 +58,10 @@ impl PyDmlStatement {
5758
Ok(self.dml.table_name.to_string())
5859
}
5960

60-
pub fn table_schema(&self) -> PyDFSchema {
61-
(*self.dml.table_schema).clone().into()
61+
pub fn target(&self) -> PyResult<PyTableSource> {
62+
Ok(PyTableSource {
63+
table_source: self.dml.target.clone(),
64+
})
6265
}
6366

6467
pub fn op(&self) -> PyWriteOp {

0 commit comments

Comments
 (0)
0