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

datafusion_python/expr/
copy_to.rs
