datafusion_python/unparser/
mod.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
18mod dialect;
19
20use std::sync::Arc;
21
22use datafusion::sql::unparser::{dialect::Dialect, Unparser};
23use dialect::PyDialect;
24use pyo3::{exceptions::PyValueError, prelude::*};
25
26use crate::sql::logical::PyLogicalPlan;
27
28#[pyclass(name = "Unparser", module = "datafusion.unparser", subclass)]
29#[derive(Clone)]
30pub struct PyUnparser {
31    dialect: Arc<dyn Dialect>,
32    pretty: bool,
33}
34
35#[pymethods]
36impl PyUnparser {
37    #[new]
38    pub fn new(dialect: PyDialect) -> Self {
39        Self {
40            dialect: dialect.dialect.clone(),
41            pretty: false,
42        }
43    }
44
45    pub fn plan_to_sql(&self, plan: &PyLogicalPlan) -> PyResult<String> {
46        let mut unparser = Unparser::new(self.dialect.as_ref());
47        unparser = unparser.with_pretty(self.pretty);
48        let sql = unparser
49            .plan_to_sql(&plan.plan())
50            .map_err(|e| PyValueError::new_err(e.to_string()))?;
51        Ok(sql.to_string())
52    }
53
54    pub fn with_pretty(&self, pretty: bool) -> Self {
55        Self {
56            dialect: self.dialect.clone(),
57            pretty,
58        }
59    }
60}
61
62pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
63    m.add_class::<PyUnparser>()?;
64    m.add_class::<PyDialect>()?;
65    Ok(())
66}
datafusion_python/unparser/mod.rs

datafusion_python/unparser/
mod.rs
