Skip to content

Commit

Permalink
Add bindings for case, cast, and trycast (apache#232)
Browse files Browse the repository at this point in the history
* Add bindings for case, cast, and trycast

* cargo clippy

* Python linters
  • Loading branch information
jdye64 authored Feb 23, 2023
1 parent f23a1a3 commit 3ca46b0
Show file tree
Hide file tree
Showing 5 changed files with 152 additions and 0 deletions.
6 changes: 6 additions & 0 deletions datafusion/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,9 @@
IsNotFalse,
IsNotUnknown,
Negative,
Case,
Cast,
TryCast,
Between,
)

Expand Down Expand Up @@ -102,6 +105,9 @@
"IsNotFalse",
"IsNotUnknown",
"Negative",
"Case",
"Cast",
"TryCast",
"Between",
]

Expand Down
6 changes: 6 additions & 0 deletions datafusion/tests/test_imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,9 @@
IsNotFalse,
IsNotUnknown,
Negative,
Case,
Cast,
TryCast,
Between,
)

Expand Down Expand Up @@ -108,6 +111,9 @@ def test_class_module_is_datafusion():
IsNotFalse,
IsNotUnknown,
Negative,
Case,
Cast,
TryCast,
Between,
]:
assert klass.__module__ == "datafusion.expr"
Expand Down
5 changes: 5 additions & 0 deletions src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@ pub mod analyze;
pub mod between;
pub mod binary_expr;
pub mod bool_expr;
pub mod case;
pub mod cast;
pub mod column;
pub mod empty_relation;
pub mod filter;
Expand Down Expand Up @@ -232,6 +234,9 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
m.add_class::<PySimilarTo>()?;
m.add_class::<PyScalarVariable>()?;
m.add_class::<alias::PyAlias>()?;
m.add_class::<case::PyCase>()?;
m.add_class::<cast::PyCast>()?;
m.add_class::<cast::PyTryCast>()?;
m.add_class::<between::PyBetween>()?;
m.add_class::<indexed_field::PyGetIndexedField>()?;
// operators
Expand Down
57 changes: 57 additions & 0 deletions src/expr/case.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
// 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 crate::expr::PyExpr;
use datafusion_expr::Case;
use pyo3::prelude::*;

#[pyclass(name = "Case", module = "datafusion.expr", subclass)]
#[derive(Clone)]
pub struct PyCase {
case: Case,
}

impl From<PyCase> for Case {
fn from(case: PyCase) -> Self {
case.case
}
}

impl From<Case> for PyCase {
fn from(case: Case) -> PyCase {
PyCase { case }
}
}

#[pymethods]
impl PyCase {
fn expr(&self) -> Option<PyExpr> {
self.case.expr.as_ref().map(|e| (**e).clone().into())
}

fn when_then_expr(&self) -> Vec<(PyExpr, PyExpr)> {
self.case
.when_then_expr
.iter()
.map(|e| ((*e.0).clone().into(), (*e.1).clone().into()))
.collect()
}

fn else_expr(&self) -> Option<PyExpr> {
self.case.else_expr.as_ref().map(|e| (**e).clone().into())
}
}
78 changes: 78 additions & 0 deletions src/expr/cast.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
// 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 crate::{common::data_type::PyDataType, expr::PyExpr};
use datafusion_expr::{Cast, TryCast};
use pyo3::prelude::*;

#[pyclass(name = "Cast", module = "datafusion.expr", subclass)]
#[derive(Clone)]
pub struct PyCast {
cast: Cast,
}

impl From<PyCast> for Cast {
fn from(cast: PyCast) -> Self {
cast.cast
}
}

impl From<Cast> for PyCast {
fn from(cast: Cast) -> PyCast {
PyCast { cast }
}
}

#[pymethods]
impl PyCast {
fn expr(&self) -> PyResult<PyExpr> {
Ok((*self.cast.expr).clone().into())
}

fn data_type(&self) -> PyResult<PyDataType> {
Ok(self.cast.data_type.clone().into())
}
}

#[pyclass(name = "TryCast", module = "datafusion.expr", subclass)]
#[derive(Clone)]
pub struct PyTryCast {
try_cast: TryCast,
}

impl From<PyTryCast> for TryCast {
fn from(try_cast: PyTryCast) -> Self {
try_cast.try_cast
}
}

impl From<TryCast> for PyTryCast {
fn from(try_cast: TryCast) -> PyTryCast {
PyTryCast { try_cast }
}
}

#[pymethods]
impl PyTryCast {
fn expr(&self) -> PyResult<PyExpr> {
Ok((*self.try_cast.expr).clone().into())
}

fn data_type(&self) -> PyResult<PyDataType> {
Ok(self.try_cast.data_type.clone().into())
}
}

0 comments on commit 3ca46b0

Please sign in to comment.