datafusion_python/expr/
copy_to.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 std::collections::HashMap;
19use std::fmt::{self, Display, Formatter};
20use std::sync::Arc;
21
22use datafusion::common::file_options::file_type::FileType;
23use datafusion::logical_expr::dml::CopyTo;
24use pyo3::prelude::*;
25use pyo3::IntoPyObjectExt;
26
27use super::logical_node::LogicalNode;
28use crate::sql::logical::PyLogicalPlan;
29
30#[pyclass(frozen, name = "CopyTo", module = "datafusion.expr", subclass)]
31#[derive(Clone)]
32pub struct PyCopyTo {
33    copy: CopyTo,
34}
35
36impl From<PyCopyTo> for CopyTo {
37    fn from(copy: PyCopyTo) -> Self {
38        copy.copy
39    }
40}
41
42impl From<CopyTo> for PyCopyTo {
43    fn from(copy: CopyTo) -> PyCopyTo {
44        PyCopyTo { copy }
45    }
46}
47
48impl Display for PyCopyTo {
49    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
50        write!(f, "CopyTo: {:?}", self.copy.output_url)
51    }
52}
53
54impl LogicalNode for PyCopyTo {
55    fn inputs(&self) -> Vec<PyLogicalPlan> {
56        vec![PyLogicalPlan::from((*self.copy.input).clone())]
57    }
58
59    fn to_variant<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
60        self.clone().into_bound_py_any(py)
61    }
62}
63
64#[pymethods]
65impl PyCopyTo {
66    #[new]
67    pub fn new(
68        input: PyLogicalPlan,
69        output_url: String,
70        partition_by: Vec<String>,
71        file_type: PyFileType,
72        options: HashMap<String, String>,
73    ) -> Self {
74        PyCopyTo {
75            copy: CopyTo::new(
76                input.plan(),
77                output_url,
78                partition_by,
79                file_type.file_type,
80                options,
81            ),
82        }
83    }
84
85    fn input(&self) -> PyLogicalPlan {
86        PyLogicalPlan::from((*self.copy.input).clone())
87    }
88
89    fn output_url(&self) -> String {
90        self.copy.output_url.clone()
91    }
92
93    fn partition_by(&self) -> Vec<String> {
94        self.copy.partition_by.clone()
95    }
96
97    fn file_type(&self) -> PyFileType {
98        PyFileType {
99            file_type: self.copy.file_type.clone(),
100        }
101    }
102
103    fn options(&self) -> HashMap<String, String> {
104        self.copy.options.clone()
105    }
106
107    fn __repr__(&self) -> PyResult<String> {
108        Ok(format!("CopyTo({self})"))
109    }
110
111    fn __name__(&self) -> PyResult<String> {
112        Ok("CopyTo".to_string())
113    }
114}
115
116#[pyclass(frozen, name = "FileType", module = "datafusion.expr", subclass)]
117#[derive(Clone)]
118pub struct PyFileType {
119    file_type: Arc<dyn FileType>,
120}
121
122impl Display for PyFileType {
123    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
124        write!(f, "FileType: {}", self.file_type)
125    }
126}
127
128#[pymethods]
129impl PyFileType {
130    fn __repr__(&self) -> PyResult<String> {
131        Ok(format!("FileType({self})"))
132    }
133
134    fn __name__(&self) -> PyResult<String> {
135        Ok("FileType".to_string())
136    }
137}
datafusion_python/expr/copy_to.rs

datafusion_python/expr/
copy_to.rs
