Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions datafusion/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@
Cast,
TryCast,
Between,
Explain,
)

__version__ = importlib_metadata.version(__name__)
Expand Down Expand Up @@ -127,6 +128,7 @@
"Cast",
"TryCast",
"Between",
"Explain",
]


Expand Down
4 changes: 4 additions & 0 deletions src/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,3 +82,7 @@ pub fn py_runtime_err(e: impl Debug) -> PyErr {
pub fn py_datafusion_err(e: impl Debug) -> PyErr {
PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("{e:?}"))
}

pub fn py_unsupported_variant_err(e: impl Debug) -> PyErr {
PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{e:?}"))
}
9 changes: 5 additions & 4 deletions src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ pub mod column;
pub mod cross_join;
pub mod empty_relation;
pub mod exists;
pub mod explain;
pub mod filter;
pub mod grouping_set;
pub mod in_list;
Expand Down Expand Up @@ -260,10 +261,7 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
m.add_class::<cast::PyTryCast>()?;
m.add_class::<between::PyBetween>()?;
m.add_class::<indexed_field::PyGetIndexedField>()?;
// operators
m.add_class::<table_scan::PyTableScan>()?;
m.add_class::<projection::PyProjection>()?;
m.add_class::<filter::PyFilter>()?;
m.add_class::<explain::PyExplain>()?;
m.add_class::<limit::PyLimit>()?;
m.add_class::<aggregate::PyAggregate>()?;
m.add_class::<sort::PySort>()?;
Expand All @@ -274,5 +272,8 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
m.add_class::<join::PyJoinConstraint>()?;
m.add_class::<cross_join::PyCrossJoin>()?;
m.add_class::<union::PyUnion>()?;
m.add_class::<filter::PyFilter>()?;
m.add_class::<projection::PyProjection>()?;
m.add_class::<table_scan::PyTableScan>()?;
Ok(())
}
6 changes: 5 additions & 1 deletion src/expr/aggregate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@ use datafusion_expr::logical_plan::Aggregate;
use pyo3::prelude::*;
use std::fmt::{self, Display, Formatter};

use super::logical_node::LogicalNode;
use crate::common::df_schema::PyDFSchema;
use crate::expr::logical_node::LogicalNode;
use crate::expr::PyExpr;
use crate::sql::logical::PyLogicalPlan;

Expand Down Expand Up @@ -103,4 +103,8 @@ impl LogicalNode for PyAggregate {
fn inputs(&self) -> Vec<PyLogicalPlan> {
vec![PyLogicalPlan::from((*self.aggregate.input).clone())]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
6 changes: 5 additions & 1 deletion src/expr/analyze.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@ use datafusion_expr::logical_plan::Analyze;
use pyo3::prelude::*;
use std::fmt::{self, Display, Formatter};

use super::logical_node::LogicalNode;
use crate::common::df_schema::PyDFSchema;
use crate::expr::logical_node::LogicalNode;
use crate::sql::logical::PyLogicalPlan;

#[pyclass(name = "Analyze", module = "datafusion.expr", subclass)]
Expand Down Expand Up @@ -77,4 +77,8 @@ impl LogicalNode for PyAnalyze {
fn inputs(&self) -> Vec<PyLogicalPlan> {
vec![PyLogicalPlan::from((*self.analyze.input).clone())]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
6 changes: 5 additions & 1 deletion src/expr/cross_join.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@ use datafusion_expr::logical_plan::CrossJoin;
use pyo3::prelude::*;
use std::fmt::{self, Display, Formatter};

use super::logical_node::LogicalNode;
use crate::common::df_schema::PyDFSchema;
use crate::expr::logical_node::LogicalNode;
use crate::sql::logical::PyLogicalPlan;

#[pyclass(name = "CrossJoin", module = "datafusion.expr", subclass)]
Expand Down Expand Up @@ -87,4 +87,8 @@ impl LogicalNode for PyCrossJoin {
PyLogicalPlan::from((*self.cross_join.right).clone()),
]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
15 changes: 14 additions & 1 deletion src/expr/empty_relation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,13 @@
// specific language governing permissions and limitations
// under the License.

use crate::common::df_schema::PyDFSchema;
use crate::{common::df_schema::PyDFSchema, sql::logical::PyLogicalPlan};
use datafusion_expr::EmptyRelation;
use pyo3::prelude::*;
use std::fmt::{self, Display, Formatter};

use super::logical_node::LogicalNode;

#[pyclass(name = "EmptyRelation", module = "datafusion.expr", subclass)]
#[derive(Clone)]
pub struct PyEmptyRelation {
Expand Down Expand Up @@ -70,3 +72,14 @@ impl PyEmptyRelation {
Ok("EmptyRelation".to_string())
}
}

impl LogicalNode for PyEmptyRelation {
fn inputs(&self) -> Vec<PyLogicalPlan> {
// table scans are leaf nodes and do not have inputs
vec![]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
110 changes: 110 additions & 0 deletions src/expr/explain.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

use std::fmt::{self, Display, Formatter};

use datafusion_expr::{logical_plan::Explain, LogicalPlan};
use pyo3::prelude::*;

use crate::{common::df_schema::PyDFSchema, errors::py_type_err, sql::logical::PyLogicalPlan};

use super::logical_node::LogicalNode;

#[pyclass(name = "Explain", module = "datafusion.expr", subclass)]
#[derive(Clone)]
pub struct PyExplain {
explain: Explain,
}

impl From<PyExplain> for Explain {
fn from(explain: PyExplain) -> Self {
explain.explain
}
}

impl From<Explain> for PyExplain {
fn from(explain: Explain) -> PyExplain {
PyExplain { explain }
}
}

impl Display for PyExplain {
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
write!(
f,
"Explain
verbose: {:?}
plan: {:?}
stringified_plans: {:?}
schema: {:?}
logical_optimization_succeeded: {:?}",
&self.explain.verbose,
&self.explain.plan,
&self.explain.stringified_plans,
&self.explain.schema,
&self.explain.logical_optimization_succeeded
)
}
}

#[pymethods]
impl PyExplain {
fn explain_string(&self) -> PyResult<Vec<String>> {
let mut string_plans: Vec<String> = Vec::new();
for stringified_plan in &self.explain.stringified_plans {
string_plans.push((*stringified_plan.plan).clone());
}
Ok(string_plans)
}

fn verbose(&self) -> bool {
self.explain.verbose
}

fn plan(&self) -> PyResult<PyLogicalPlan> {
Ok(PyLogicalPlan::from((*self.explain.plan).clone()))
}

fn schema(&self) -> PyDFSchema {
(*self.explain.schema).clone().into()
}

fn logical_optimization_succceeded(&self) -> bool {
self.explain.logical_optimization_succeeded
}
}

impl TryFrom<LogicalPlan> for PyExplain {
type Error = PyErr;

fn try_from(logical_plan: LogicalPlan) -> Result<Self, Self::Error> {
match logical_plan {
LogicalPlan::Explain(explain) => Ok(PyExplain { explain }),
_ => Err(py_type_err("unexpected plan")),
}
}
}

impl LogicalNode for PyExplain {
fn inputs(&self) -> Vec<PyLogicalPlan> {
vec![]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
8 changes: 6 additions & 2 deletions src/expr/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@ impl Display for PyFilter {
write!(
f,
"Filter
\nPredicate: {:?}
\nInput: {:?}",
Predicate: {:?}
Input: {:?}",
&self.filter.predicate, &self.filter.input
)
}
Expand Down Expand Up @@ -80,4 +80,8 @@ impl LogicalNode for PyFilter {
fn inputs(&self) -> Vec<PyLogicalPlan> {
vec![PyLogicalPlan::from((*self.filter.input).clone())]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
20 changes: 12 additions & 8 deletions src/expr/join.rs
Original file line number Diff line number Diff line change
Expand Up @@ -95,14 +95,14 @@ impl Display for PyJoin {
write!(
f,
"Join
\nLeft: {:?}
\nRight: {:?}
\nOn: {:?}
\nFilter: {:?}
\nJoinType: {:?}
\nJoinConstraint: {:?}
\nSchema: {:?}
\nNullEqualsNull: {:?}",
Left: {:?}
Right: {:?}
On: {:?}
Filter: {:?}
JoinType: {:?}
JoinConstraint: {:?}
Schema: {:?}
NullEqualsNull: {:?}",
&self.join.left,
&self.join.right,
&self.join.on,
Expand Down Expand Up @@ -178,4 +178,8 @@ impl LogicalNode for PyJoin {
PyLogicalPlan::from((*self.join.right).clone()),
]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
10 changes: 7 additions & 3 deletions src/expr/limit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,9 @@ impl Display for PyLimit {
write!(
f,
"Limit
\nSkip: {}
\nFetch: {:?}
\nInput: {:?}",
Skip: {}
Fetch: {:?}
Input: {:?}",
&self.limit.skip, &self.limit.fetch, &self.limit.input
)
}
Expand Down Expand Up @@ -85,4 +85,8 @@ impl LogicalNode for PyLimit {
fn inputs(&self) -> Vec<PyLogicalPlan> {
vec![PyLogicalPlan::from((*self.limit.input).clone())]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
4 changes: 4 additions & 0 deletions src/expr/logical_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,11 +15,15 @@
// specific language governing permissions and limitations
// under the License.

use pyo3::{PyObject, PyResult, Python};

use crate::sql::logical::PyLogicalPlan;

/// Representation of a `LogicalNode` in the in overall `LogicalPlan`
/// any "node" shares these common traits in common.
pub trait LogicalNode {
/// The input plan to the current logical node instance.
fn inputs(&self) -> Vec<PyLogicalPlan>;

fn to_variant(&self, py: Python) -> PyResult<PyObject>;
}
4 changes: 4 additions & 0 deletions src/expr/projection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -96,4 +96,8 @@ impl LogicalNode for PyProjection {
fn inputs(&self) -> Vec<PyLogicalPlan> {
vec![PyLogicalPlan::from((*self.projection.input).clone())]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
4 changes: 4 additions & 0 deletions src/expr/sort.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,4 +91,8 @@ impl LogicalNode for PySort {
fn inputs(&self) -> Vec<PyLogicalPlan> {
vec![PyLogicalPlan::from((*self.sort.input).clone())]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
10 changes: 7 additions & 3 deletions src/expr/table_scan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,9 +52,9 @@ impl Display for PyTableScan {
write!(
f,
"TableScan\nTable Name: {}
\nProjections: {:?}
\nProjected Schema: {:?}
\nFilters: {:?}",
Projections: {:?}
Projected Schema: {:?}
Filters: {:?}",
&self.table_scan.table_name,
&self.py_projections(),
&self.py_schema(),
Expand Down Expand Up @@ -131,4 +131,8 @@ impl LogicalNode for PyTableScan {
// table scans are leaf nodes and do not have inputs
vec![]
}

fn to_variant(&self, py: Python) -> PyResult<PyObject> {
Ok(self.clone().into_py(py))
}
}
Loading