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

Commit afd28a6

Browse files
chenkovskychenkovsky
committed
add expr
1 parent 40a61c1 commit afd28a6

17 files changed

+1885
-13
lines changed

‎src/expr.rs

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,21 @@ pub mod union;
9696
pub mod unnest;
9797
pub mod unnest_expr;
9898
pub mod window;
99+
pub mod statement;
100+
pub mod values;
101+
pub mod dml;
102+
pub mod create_external_table;
103+
pub mod copy_to;
104+
pub mod create_catalog_schema;
105+
pub mod drop_view;
106+
pub mod create_catalog;
107+
pub mod drop_catalog_schema;
108+
pub mod drop_function;
109+
pub mod create_function;
110+
pub mod create_index;
111+
pub mod describe_table;
112+
pub mod recursive_query;
113+
pub mod constraints;
99114

100115
use sort_expr::{to_sort_expressions, PySortExpr};
101116

@@ -784,5 +799,33 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
784799
m.add_class::<window::PyWindowExpr>()?;
785800
m.add_class::<window::PyWindowFrame>()?;
786801
m.add_class::<window::PyWindowFrameBound>()?;
802+
m.add_class::<copy_to::PyCopyTo>()?;
803+
m.add_class::<copy_to::PyFileType>()?;
804+
m.add_class::<create_catalog::PyCreateCatalog>()?;
805+
m.add_class::<create_catalog_schema::PyCreateCatalogSchema>()?;
806+
m.add_class::<create_external_table::PyCreateExternalTable>()?;
807+
m.add_class::<constraints::PyConstraints>()?;
808+
m.add_class::<create_function::PyCreateFunction>()?;
809+
m.add_class::<create_function::PyOperateFunctionArg>()?;
810+
m.add_class::<create_function::PyCreateFunctionBody>()?;
811+
m.add_class::<create_index::PyCreateIndex>()?;
812+
m.add_class::<describe_table::PyDescribeTable>()?;
813+
m.add_class::<dml::PyDmlStatement>()?;
814+
m.add_class::<drop_catalog_schema::PyDropCatalogSchema>()?;
815+
m.add_class::<drop_function::PyDropFunction>()?;
816+
m.add_class::<drop_view::PyDropView>()?;
817+
m.add_class::<recursive_query::PyRecursiveQuery>()?;
818+
819+
m.add_class::<statement::PyTransactionStart>()?;
820+
m.add_class::<statement::PyTransactionEnd>()?;
821+
m.add_class::<statement::PySetVariable>()?;
822+
m.add_class::<statement::PyPrepare>()?;
823+
m.add_class::<statement::PyExecute>()?;
824+
m.add_class::<statement::PyDeallocate>()?;
825+
m.add_class::<values::PyValues>()?;
826+
m.add_class::<statement::PyTransactionAccessMode>()?;
827+
m.add_class::<statement::PyTransactionConclusion>()?;
828+
m.add_class::<statement::PyTransactionIsolationLevel>()?;
829+
787830
Ok(())
788831
}

‎src/expr/constraints.rs

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
use std::fmt::{self, Display, Formatter};
2+
3+
use datafusion::common::Constraints;
4+
use pyo3::prelude::*;
5+
6+
#[pyclass(name = "Constraints", module = "datafusion.expr", subclass)]
7+
#[derive(Clone)]
8+
pub struct PyConstraints {
9+
pub constraints: Constraints,
10+
}
11+
12+
impl From<PyConstraints> for Constraints {
13+
fn from(constraints: PyConstraints) -> Self {
14+
constraints.constraints
15+
}
16+
}
17+
18+
impl From<Constraints> for PyConstraints {
19+
fn from(constraints: Constraints) -> Self {
20+
PyConstraints { constraints }
21+
}
22+
}
23+
24+
impl Display for PyConstraints {
25+
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
26+
write!(f, "Constraints: {:?}", self.constraints)
27+
}
28+
}
29+

‎src/expr/copy_to.rs

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
use std::{collections::HashMap, fmt::{self, Display, Formatter}, sync::Arc};
19+
20+
use datafusion::{common::file_options::file_type::FileType, logical_expr::dml::CopyTo};
21+
use pyo3::prelude::*;
22+
23+
use crate::sql::logical::PyLogicalPlan;
24+
25+
use super::logical_node::LogicalNode;
26+
27+
#[pyclass(name = "CopyTo", module = "datafusion.expr", subclass)]
28+
#[derive(Clone)]
29+
pub struct PyCopyTo {
30+
copy: CopyTo,
31+
}
32+
33+
impl From<PyCopyTo> for CopyTo {
34+
fn from(copy: PyCopyTo) -> Self {
35+
copy.copy
36+
}
37+
}
38+
39+
impl From<CopyTo> for PyCopyTo {
40+
fn from(copy: CopyTo) -> PyCopyTo {
41+
PyCopyTo { copy }
42+
}
43+
}
44+
45+
impl Display for PyCopyTo {
46+
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
47+
write!(f, "CopyTo: {:?}", self.copy.output_url)
48+
}
49+
}
50+
51+
impl LogicalNode for PyCopyTo {
52+
53+
fn inputs(&self) -> Vec<PyLogicalPlan> {
54+
vec![PyLogicalPlan::from((*self.copy.input).clone())]
55+
}
56+
57+
fn to_variant(&self, py: Python) -> PyResult<PyObject> {
58+
Ok(self.clone().into_py(py))
59+
}
60+
}
61+
62+
#[pymethods]
63+
impl PyCopyTo {
64+
#[new]
65+
pub fn new(
66+
input: PyLogicalPlan,
67+
output_url: String,
68+
partition_by: Vec<String>,
69+
file_type: PyFileType,
70+
options: HashMap<String, String>) -> Self {
71+
return PyCopyTo {
72+
copy: CopyTo {
73+
input: input.plan(),
74+
output_url,
75+
partition_by,
76+
file_type: file_type.file_type,
77+
options,
78+
},
79+
}
80+
}
81+
82+
fn input(&self) -> PyLogicalPlan {
83+
PyLogicalPlan::from((*self.copy.input).clone())
84+
}
85+
86+
fn output_url(&self) -> String {
87+
self.copy.output_url.clone()
88+
}
89+
90+
fn partition_by(&self) -> Vec<String> {
91+
self.copy.partition_by.clone()
92+
}
93+
94+
fn file_type(&self) -> PyFileType {
95+
PyFileType { file_type: self.copy.file_type.clone() }
96+
}
97+
98+
fn options(&self) -> HashMap<String, String> {
99+
self.copy.options.clone()
100+
}
101+
102+
fn __repr__(&self) -> PyResult<String> {
103+
Ok(format!("CopyTo({})", self))
104+
}
105+
106+
fn __name__(&self) -> PyResult<String> {
107+
Ok("CopyTo".to_string())
108+
}
109+
}
110+
111+
#[pyclass(name = "FileType", module = "datafusion.expr", subclass)]
112+
#[derive(Clone)]
113+
pub struct PyFileType {
114+
file_type: Arc<dyn FileType>,
115+
}
116+
117+
impl Display for PyFileType {
118+
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
119+
write!(f, "FileType: {}", self.file_type)
120+
}
121+
}
122+
123+
#[pymethods]
124+
impl PyFileType {
125+
fn __repr__(&self) -> PyResult<String> {
126+
Ok(format!("FileType({})", self))
127+
}
128+
129+
fn __name__(&self) -> PyResult<String> {
130+
Ok("FileType".to_string())
131+
}
132+
}

‎src/expr/create_catalog.rs

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
use std::{fmt::{self, Display, Formatter}, sync::Arc};
19+
20+
use datafusion::logical_expr::CreateCatalog;
21+
use pyo3::prelude::*;
22+
23+
use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
24+
25+
use super::logical_node::LogicalNode;
26+
27+
#[pyclass(name = "CreateCatalog", module = "datafusion.expr", subclass)]
28+
#[derive(Clone)]
29+
pub struct PyCreateCatalog {
30+
create: CreateCatalog,
31+
}
32+
33+
impl From<PyCreateCatalog> for CreateCatalog {
34+
fn from(create: PyCreateCatalog) -> Self {
35+
create.create
36+
}
37+
}
38+
39+
impl From<CreateCatalog> for PyCreateCatalog {
40+
fn from(create: CreateCatalog) -> PyCreateCatalog {
41+
PyCreateCatalog { create }
42+
}
43+
}
44+
45+
impl Display for PyCreateCatalog {
46+
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
47+
write!(f, "CreateCatalog: {:?}", self.create.catalog_name)
48+
}
49+
}
50+
51+
#[pymethods]
52+
impl PyCreateCatalog {
53+
54+
#[new]
55+
pub fn new(catalog_name: String, if_not_exists: bool, schema: PyDFSchema) -> PyResult<PyCreateCatalog> {
56+
Ok(PyCreateCatalog {
57+
create: CreateCatalog {
58+
catalog_name,
59+
if_not_exists,
60+
schema: Arc::new(schema.into()),
61+
}
62+
})
63+
}
64+
65+
pub fn catalog_name(&self) -> String {
66+
self.create.catalog_name.clone()
67+
}
68+
69+
pub fn if_not_exists(&self) -> bool {
70+
self.create.if_not_exists
71+
}
72+
73+
pub fn schema(&self) -> PyDFSchema {
74+
(*self.create.schema).clone().into()
75+
}
76+
77+
fn __repr__(&self) -> PyResult<String> {
78+
Ok(format!("CreateCatalog({})", self))
79+
}
80+
81+
fn __name__(&self) -> PyResult<String> {
82+
Ok("CreateCatalog".to_string())
83+
}
84+
}
85+
86+
impl LogicalNode for PyCreateCatalog {
87+
fn inputs(&self) -> Vec<PyLogicalPlan> {
88+
vec![]
89+
}
90+
91+
fn to_variant(&self, py: Python) -> PyResult<PyObject> {
92+
Ok(self.clone().into_py(py))
93+
}
94+
}

0 commit comments

Comments
3148  (0)
0