add expr

Author: chenkovsky

src/expr.rs

@@ -96,6 +96,21 @@ pub mod union;
 pub mod unnest;
 pub mod unnest_expr;
 pub mod window;
+pub mod statement;
+pub mod values;
+pub mod dml;
+pub mod create_external_table;
+pub mod copy_to;
+pub mod create_catalog_schema;
+pub mod drop_view;
+pub mod create_catalog;
+pub mod drop_catalog_schema;
+pub mod drop_function;
+pub mod create_function;
+pub mod create_index;
+pub mod describe_table;
+pub mod recursive_query;
+pub mod constraints;
 
 use sort_expr::{to_sort_expressions, PySortExpr};
 
@@ -784,5 +799,33 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<window::PyWindowExpr>()?;
     m.add_class::<window::PyWindowFrame>()?;
     m.add_class::<window::PyWindowFrameBound>()?;
+    m.add_class::<copy_to::PyCopyTo>()?;
+    m.add_class::<copy_to::PyFileType>()?;
+    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>()?;
+    m.add_class::<create_index::PyCreateIndex>()?;
+    m.add_class::<describe_table::PyDescribeTable>()?;
+    m.add_class::<dml::PyDmlStatement>()?;
+    m.add_class::<drop_catalog_schema::PyDropCatalogSchema>()?;
+    m.add_class::<drop_function::PyDropFunction>()?;
+    m.add_class::<drop_view::PyDropView>()?;
+    m.add_class::<recursive_query::PyRecursiveQuery>()?;
+
+    m.add_class::<statement::PyTransactionStart>()?;
+    m.add_class::<statement::PyTransactionEnd>()?;
+    m.add_class::<statement::PySetVariable>()?;
+    m.add_class::<statement::PyPrepare>()?;
+    m.add_class::<statement::PyExecute>()?;
+    m.add_class::<statement::PyDeallocate>()?;
+    m.add_class::<values::PyValues>()?;
+    m.add_class::<statement::PyTransactionAccessMode>()?;
+    m.add_class::<statement::PyTransactionConclusion>()?;
+    m.add_class::<statement::PyTransactionIsolationLevel>()?;
+
     Ok(())
 }

src/expr/constraints.rs

@@ -0,0 +1,29 @@
+use std::fmt::{self, Display, Formatter};
+
+use datafusion::common::Constraints;
+use pyo3::prelude::*;
+
+#[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)
+    }
+}
+

src/expr/copy_to.rs

@@ -0,0 +1,132 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::{collections::HashMap, fmt::{self, Display, Formatter}, sync::Arc};
+
+use datafusion::{common::file_options::file_type::FileType, logical_expr::dml::CopyTo};
+use pyo3::prelude::*;
+
+use crate::sql::logical::PyLogicalPlan;
+
+use super::logical_node::LogicalNode;
+
+#[pyclass(name = "CopyTo", module = "datafusion.expr", subclass)]
+#[derive(Clone)]
+pub struct PyCopyTo {
+    copy: CopyTo,
+}
+
+impl From<PyCopyTo> for CopyTo {
+    fn from(copy: PyCopyTo) -> Self {
+        copy.copy
+    }
+}
+
+impl From<CopyTo> for PyCopyTo {
+    fn from(copy: CopyTo) -> PyCopyTo {
+        PyCopyTo { copy }
+    }
+}
+
+impl Display for PyCopyTo {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        write!(f, "CopyTo: {:?}", self.copy.output_url)
+    }
+}
+
+impl LogicalNode for PyCopyTo {
+
+    fn inputs(&self) -> Vec<PyLogicalPlan> {
+        vec![PyLogicalPlan::from((*self.copy.input).clone())]
+    }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
+}
+
+#[pymethods]
+impl PyCopyTo {
+    #[new]
+    pub fn new(
+        input: PyLogicalPlan,
+        output_url: String,
+        partition_by: Vec<String>,
+        file_type: PyFileType,
+        options: HashMap<String, String>) -> Self {
+        return PyCopyTo {
+            copy: CopyTo {
+                input: input.plan(),
+                output_url,
+                partition_by,
+                file_type: file_type.file_type,
+                options,
+            },
+        }
+    }
+
+    fn input(&self) -> PyLogicalPlan {
+        PyLogicalPlan::from((*self.copy.input).clone())
+    }
+
+    fn output_url(&self) -> String {
+        self.copy.output_url.clone()
+    }
+
+    fn partition_by(&self) -> Vec<String> {
+        self.copy.partition_by.clone()
+    }
+
+    fn file_type(&self) -> PyFileType {
+        PyFileType { file_type: self.copy.file_type.clone() }
+    }
+
+    fn options(&self) -> HashMap<String, String> {
+        self.copy.options.clone()
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!("CopyTo({})", self))
+    }
+
+    fn __name__(&self) -> PyResult<String> {
+        Ok("CopyTo".to_string())
+    }
+}
+
+#[pyclass(name = "FileType", module = "datafusion.expr", subclass)]
+#[derive(Clone)]
+pub struct PyFileType {
+    file_type: Arc<dyn FileType>,
+}
+
+impl Display for PyFileType {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        write!(f, "FileType: {}", self.file_type)
+    }
+}
+
+#[pymethods]
+impl PyFileType {
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!("FileType({})", self))
+    }
+
+    fn __name__(&self) -> PyResult<String> {
+        Ok("FileType".to_string())
+    }
+}

src/expr/create_catalog.rs

@@ -0,0 +1,94 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::{fmt::{self, Display, Formatter}, sync::Arc};
+
+use datafusion::logical_expr::CreateCatalog;
+use pyo3::prelude::*;
+
+use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
+
+use super::logical_node::LogicalNode;
+
+#[pyclass(name = "CreateCatalog", module = "datafusion.expr", subclass)]
+#[derive(Clone)]
+pub struct PyCreateCatalog {
+    create: CreateCatalog,
+}
+
+impl From<PyCreateCatalog> for CreateCatalog {
+    fn from(create: PyCreateCatalog) -> Self {
+        create.create
+    }
+}
+
+impl From<CreateCatalog> for PyCreateCatalog {
+    fn from(create: CreateCatalog) -> PyCreateCatalog {
+        PyCreateCatalog { create }
+    }
+}
+
+impl Display for PyCreateCatalog {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        write!(f, "CreateCatalog: {:?}", self.create.catalog_name)
+    }
+}
+
+#[pymethods]
+impl PyCreateCatalog {
+
+    #[new]
+    pub fn new(catalog_name: String, if_not_exists: bool, schema: PyDFSchema) -> PyResult<PyCreateCatalog> {
+        Ok(PyCreateCatalog {
+            create: CreateCatalog {
+                catalog_name,
+                if_not_exists,
+                schema: Arc::new(schema.into()),
+            }
+        })
+    }
+
+    pub fn catalog_name(&self) -> String {
+        self.create.catalog_name.clone()
+    }
+
+    pub fn if_not_exists(&self) -> bool {
+        self.create.if_not_exists
+    }
+
+    pub fn schema(&self) -> PyDFSchema {
+        (*self.create.schema).clone().into()
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!("CreateCatalog({})", self))
+    }
+
+    fn __name__(&self) -> PyResult<String> {
+        Ok("CreateCatalog".to_string())
+    }
+}
+
+impl LogicalNode for PyCreateCatalog {
+    fn inputs(&self) -> Vec<PyLogicalPlan> {
+        vec![]
+    }
+
+    fn to_variant(&self, py: Python) -> PyResult<PyObject> {
+        Ok(self.clone().into_py(py))
+    }
+}