8000 feat: support unparser (#1088) · satwikmishra11/datafusion-python@ffafb59 · GitHub
[go: up one dir, main page]
Skip to content

Commit ffafb59

Browse files
chenkovskychenkovsky
authored
feat: support unparser (apache#1088)
* support unparser * add license * add export * format * format
1 parent 583e1e9 commit ffafb59

File tree

6 files changed

+249
-1
lines changed

6 files changed

+249
-1
lines changed

‎python/datafusion/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
except ImportError:
2727
import importlib_metadata
2828

29-
from . import functions, object_store, substrait
29+
from . import functions, object_store, substrait, unparser
3030

3131
# The following imports are okay to remain as opaque to the user.
3232
from ._internal import Config
@@ -89,6 +89,7 @@
8989
"udaf",
9090
"udf",
9191
"udwf",
92+
"unparser",
9293
]
9394

9495

‎python/datafusion/unparser.py

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
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+
18+
"""This module provides support for unparsing datafusion plans to SQL.
19+
20+
For additional information about unparsing, see https://docs.rs/datafusion-sql/latest/datafusion_sql/unparser/index.html
21+
"""
22+
23+
from ._internal import unparser as unparser_internal
24+
from .plan import LogicalPlan
25+
26+
27+
class Dialect:
28+
"""DataFusion data catalog."""
29+
30+
def __init__(self, dialect: unparser_internal.Dialect) -> None:
31+
"""This constructor is not typically called by the end user."""
32+
self.dialect = dialect
33+
34+
@staticmethod
35+
def default() -> "Dialect":
36+
"""Create a new default dialect."""
37+
return Dialect(unparser_internal.Dialect.default())
38+
39+
@staticmethod
40+
def mysql() -> "Dialect":
41+
"""Create a new MySQL dialect."""
42+
return Dialect(unparser_internal.Dialect.mysql())
43+
44+
@staticmethod
45+
def postgres() -> "Dialect":
46+
"""Create a new PostgreSQL dialect."""
47+
return Dialect(unparser_internal.Dialect.postgres())
48+
49+
@staticmethod
50+
def sqlite() -> "Dialect":
51+
"""Create a new SQLite dialect."""
52+
return Dialect(unparser_internal.Dialect.sqlite())
53+
54+
@staticmethod
55+
def duckdb() -> "Dialect":
56+
"""Create a new DuckDB dialect."""
57+
return Dialect(unparser_internal.Dialect.duckdb())
58+
59+
60+
class Unparser:
61+
"""DataFusion unparser."""
62+
63+
def __init__(self, dialect: Dialect) -> None:
64+
"""This constructor is not typically called by the end user."""
65+
self.unparser = unparser_internal.Unparser(dialect.dialect)
66+
67+
def plan_to_sql(self, plan: LogicalPlan) -> str:
68+
"""Convert a logical plan to a SQL string."""
69+
return self.unparser.plan_to_sql(plan._raw_plan)
70+
71+
def with_pretty(self, pretty: bool) -> "Unparser":
72+
"""Set the pretty flag."""
73+
self.unparser = self.unparser.with_pretty(pretty)
74+
return self
75+
76+
77+
__all__ = [
78+
"Dialect",
79+
"Unparser",
80+
]

‎python/tests/test_unparser.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
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+
18+
from datafusion.context import SessionContext
19+
from datafusion.unparser import Dialect, Unparser
20+
21+
22+
def test_unparser():
23+
ctx = SessionContext()
24+
df = ctx.sql("SELECT 1")
25+
for dialect in [
26+
Dialect.mysql(),
27+
Dialect.postgres(),
28+
Dialect.sqlite(),
29+
Dialect.duckdb(),
30+
]:
31+
unparser = Unparser(dialect)
32+
sql = unparser.plan_to_sql(df.logical_plan())
33+
assert sql == "SELECT 1"

‎src/lib.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ pub mod pyarrow_util;
5252
mod record_batch;
5353
pub mod sql;
5454
pub mod store;
55+
pub mod unparser;
5556

5657
#[cfg(feature = "substrait")]
5758
pub mod substrait;
@@ -103,6 +104,10 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
103104
expr::init_module(&expr)?;
104105
m.add_submodule(&expr)?;
105106

107+
let unparser = PyModule::new(py, "unparser")?;
108+
unparser::init_module(&unparser)?;
109+
m.add_submodule(&unparser)?;
110+
106111
// Register the functions as a submodule
107112
let funcs = PyModule::new(py, "functions")?;
108113
functions::init_module(&funcs)?;

‎src/unparser/dialect.rs

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
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+
18+
use std::sync::Arc;
19+
20+
use datafusion::sql::unparser::dialect::{
21+
DefaultDialect, Dialect, DuckDBDialect, MySqlDialect, PostgreSqlDialect, SqliteDialect,
22+
};
23+
use pyo3::prelude::*;
24+
25+
#[pyclass(name = "Dialect", module = "datafusion.unparser", subclass)]
26+
#[derive(Clone)]
27+
pub struct PyDialect {
28+
pub dialect: Arc<dyn Dialect>,
29+
}
30+
31+
#[pymethods]
32+
impl PyDialect {
33+
#[staticmethod]
34+
pub fn default() -> Self {
35+
Self {
36+
dialect: Arc::new(DefaultDialect {}),
37+
}
38+
}
39+
#[staticmethod]
40+
pub fn postgres() -> Self {
41+
Self {
42+
dialect: Arc::new(PostgreSqlDialect {}),
43+
}
44+
}
45+
#[staticmethod]
46+
pub fn mysql() -> Self {
47+
Self {
48+
dialect: Arc::new(MySqlDialect {}),
49+
}
50+
}
51+
#[staticmethod]
52+
pub fn sqlite() -> Self {
53+
Self {
54+
dialect: Arc::new(SqliteDialect {}),
55+
}
56+
}
57+
#[staticmethod]
58+
pub fn duckdb() -> Self {
59+
Self {
60+
dialect: Arc::new(DuckDBDialect::new()),
61+
}
62+
}
63+
}

‎src/unparser/mod.rs

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
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+
18+
mod dialect;
19+
20+
use std::sync::Arc;
21+
22+
use datafusion::sql::unparser::{dialect::Dialect, Unparser};
23+
use dialect::PyDialect;
24+
use pyo3::{exceptions::PyValueError, prelude::*};
25+
26+
use crate::sql::logical::PyLogicalPlan;
27+
28+
#[pyclass(name = "Unparser", module = "datafusion.unparser", subclass)]
29+
#[derive(Clone)]
30+
pub struct PyUnparser {
31+
dialect: Arc<dyn Dialect>,
32+
pretty: bool,
33+
}
34+
35+
#[pymethods]
36+
impl 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+
62+
pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
63+
m.add_class::<PyUnparser>()?;
64+
m.add_class::<PyDialect>()?;
65+
Ok(())
66+
}

0 commit comments

Comments
 (0)
0