datafusion_python/expr/
dml.rs

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
18use datafusion::logical_expr::dml::InsertOp;
19use datafusion::logical_expr::{DmlStatement, WriteOp};
20use pyo3::IntoPyObjectExt;
21use pyo3::prelude::*;
22
23use super::logical_node::LogicalNode;
24use crate::common::df_schema::PyDFSchema;
25use crate::common::schema::PyTableSource;
26use crate::sql::logical::PyLogicalPlan;
27
28#[pyclass(
29    from_py_object,
30    frozen,
31    name = "DmlStatement",
32    module = "datafusion.expr",
33    subclass
34)]
35#[derive(Clone)]
36pub struct PyDmlStatement {
37    dml: DmlStatement,
38}
39
40impl From<PyDmlStatement> for DmlStatement {
41    fn from(dml: PyDmlStatement) -> Self {
42        dml.dml
43    }
44}
45
46impl From<DmlStatement> for PyDmlStatement {
47    fn from(dml: DmlStatement) -> PyDmlStatement {
48        PyDmlStatement { dml }
49    }
50}
51
52impl LogicalNode for PyDmlStatement {
53    fn inputs(&self) -> Vec<PyLogicalPlan> {
54        vec![PyLogicalPlan::from((*self.dml.input).clone())]
55    }
56
57    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
58        self.clone().into_bound_py_any(py)
59    }
60}
61
62#[pymethods]
63impl PyDmlStatement {
64    pub fn table_name(&self) -> PyResult<String> {
65        Ok(self.dml.table_name.to_string())
66    }
67
68    pub fn target(&self) -> PyResult<PyTableSource> {
69        Ok(PyTableSource {
70            table_source: self.dml.target.clone(),
71        })
72    }
73
74    pub fn op(&self) -> PyWriteOp {
75        self.dml.op.clone().into()
76    }
77
78    pub fn input(&self) -> PyLogicalPlan {
79        PyLogicalPlan {
80            plan: self.dml.input.clone(),
81        }
82    }
83
84    pub fn output_schema(&self) -> PyDFSchema {
85        (*self.dml.output_schema).clone().into()
86    }
87
88    fn __repr__(&self) -> PyResult<String> {
89        Ok("DmlStatement".to_string())
90    }
91
92    fn __name__(&self) -> PyResult<String> {
93        Ok("DmlStatement".to_string())
94    }
95}
96
97#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
98#[pyclass(
99    from_py_object,
100    eq,
101    eq_int,
102    name = "WriteOp",
103    module = "datafusion.expr"
104)]
105pub enum PyWriteOp {
106    Append,
107    Overwrite,
108    Replace,
109    Update,
110    Delete,
111    Ctas,
112    Truncate,
113}
114
115impl From<WriteOp> for PyWriteOp {
116    fn from(write_op: WriteOp) -> Self {
117        match write_op {
118            WriteOp::Insert(InsertOp::Append) => PyWriteOp::Append,
119            WriteOp::Insert(InsertOp::Overwrite) => PyWriteOp::Overwrite,
120            WriteOp::Insert(InsertOp::Replace) => PyWriteOp::Replace,
121            WriteOp::Update => PyWriteOp::Update,
122            WriteOp::Delete => PyWriteOp::Delete,
123            WriteOp::Ctas => PyWriteOp::Ctas,
124            WriteOp::Truncate => PyWriteOp::Truncate,
125        }
126    }
127}
128
129impl From<PyWriteOp> for WriteOp {
130    fn from(py: PyWriteOp) -> Self {
131        match py {
132            PyWriteOp::Append => WriteOp::Insert(InsertOp::Append),
133            PyWriteOp::Overwrite => WriteOp::Insert(InsertOp::Overwrite),
134            PyWriteOp::Replace => WriteOp::Insert(InsertOp::Replace),
135            PyWriteOp::Update => WriteOp::Update,
136            PyWriteOp::Delete => WriteOp::Delete,
137            PyWriteOp::Ctas => WriteOp::Ctas,
138            PyWriteOp::Truncate => WriteOp::Truncate,
139        }
140    }
141}
142
143#[pymethods]
144impl PyWriteOp {
145    fn name(&self) -> String {
146        let write_op: WriteOp = self.clone().into();
147        write_op.name().to_string()
148    }
149}
datafusion_python/expr/dml.rs

datafusion_python/expr/
dml.rs
