8000 Add bindings for case, cast, and trycast (#232) · chenqin/arrow-datafusion-python@3ca46b0 · GitHub
[go: up one dir, main page]
Skip to content

Commit 3ca46b0

Browse files
jdye64jdye64
authored
Add bindings for case, cast, and trycast (apache#232)
* Add bindings for case, cast, and trycast * cargo clippy * Python linters
1 parent f23a1a3 commit 3ca46b0

File tree

5 files changed

+152
-0
lines changed

5 files changed

+152
-0
lines changed

‎datafusion/__init__.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,9 @@
6363
IsNotFalse,
6464
IsNotUnknown,
6565
Negative,
66+
Case,
67+
Cast,
68+
TryCast,
6669
Between,
6770
)
6871

@@ -102,6 +105,9 @@
102105
"IsNotFalse",
103106
"IsNotUnknown",
104107
"Negative",
108+
"Case",
109+
"Cast",
110+
"TryCast",
105111
"Between",
106112
]
107113

‎datafusion/tests/test_imports.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,9 @@
5959
IsNotFalse,
6060
IsNotUnknown,
6161
Negative,
62+
Case,
63+
Cast,
64+
TryCast,
6265
Between,
6366
)
6467

@@ -108,6 +111,9 @@ def test_class_module_is_datafusion():
108111
IsNotFalse,
109112
IsNotUnknown,
110113
Negative,
114+
Case,
115+
Cast,
116+
TryCast,
111117
Between,
112118
]:
113119
assert klass.__module__ == "datafusion.expr"

‎src/expr.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ pub mod analyze;
4444
pub mod between;
4545
pub mod binary_expr;
4646
pub mod bool_expr;
47+
pub mod case;
48+
pub mod cast;
4749
pub mod column;
4850
pub mod empty_relation;
4951
pub mod filter;
@@ -232,6 +234,9 @@ pub(crate) fn init_module(m: &PyModule) -> PyResult<()> {
232234
m.add_class::<PySimilarTo>()?;
233235
m.add_class::<PyScalarVariable>()?;
234236
m.add_class::<alias::PyAlias>()?;
237+
m.add_class::<case::PyCase>()?;
238+
m.add_class::<cast::PyCast>()?;
239+
m.add_class::<cast::PyTryCast>()?;
235240
m.add_class::<between::PyBetween>()?;
236241
m.add_class::<indexed_field::PyGetIndexedField>()?;
237242
// operators

‎src/expr/case.rs

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
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 crate::expr::PyExpr;
19+
use datafusion_expr::Case;
20+
use pyo3::prelude::*;
21+
22+
#[pyclass(name = "Case", module = "datafusion.expr", subclass)]
23+
#[derive(Clone)]
24+
pub struct PyCase {
25+
case: Case,
26+
}
27+
28+
impl From<PyCase> for Case {
29+
fn from(case: PyCase) -> Self {
30+
case.case
31+
}
32+
}
33+
34+
impl From<Case> for PyCase {
35+
fn from(case: Case) -> PyCase {
36+
PyCase { case }
37+
}
38+
}
39+
40+
#[pymethods]
41+
impl PyCase {
42+
fn expr(&self) -> Option<PyExpr> {
43+
self.case.expr.as_ref().map(|e| (**e).clone().into())
44+
}
45+
46+
fn when_then_expr(&self) -> Vec<(PyExpr, PyExpr)> {
47+
self.case
48+
.when_then_expr
49+
.iter()
50+
.map(|e| ((*e.0).clone().into(), (*e.1).clone().into()))
51+
.collect()
52+
}
53+
54+
fn else_expr(&self) -> Option<PyExpr> {
55+
self.case.e F438 lse_expr.as_ref().map(|e| (**e).clone().into())
56+
}
57+
}

‎src/expr/cast.rs

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
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 crate::{common::data_type::PyDataType, expr::PyExpr};
19+
use datafusion_expr::{Cast, TryCast};
20+
use pyo3::prelude::*;
21+
22+
#[pyclass(name = "Cast", module = "datafusion.expr", subclass)]
23+
#[derive(Clone)]
24+
pub struct PyCast {
25+
cast: Cast,
26+
}
27+
28+
impl From<PyCast> for Cast {
29+
fn from(cast: PyCast) -> Self {
30+
cast.cast
31+
}
32+
}
33+
34+
impl From<Cast> for PyCast {
35+
fn from(cast: Cast) -> PyCast {
36+
PyCast { cast }
37+
}
38+
}
39+
40+
#[pymethods]
41+
impl PyCast {
42+
fn expr(&self) -> PyResult<PyExpr> {
43+
Ok((*self.cast.expr).clone().into())
44+
}
45+
46+
fn data_type(&self) -> PyResult<PyDataType> {
47+
Ok(self.cast.data_type.clone().into())
48+
}
49+
}
50+
51+
#[pyclass(name = "TryCast", module = "datafusion.expr", subclass)]
52+
#[derive(Clone)]
53+
pub struct PyTryCast {
54+
try_cast: TryCast,
55+
}
56+
57+
impl From<PyTryCast> for TryCast {
58+
fn from(try_cast: PyTryCast) -> Self {
59+
try_cast.try_cast
60+
}
61+
}
62+
63+
impl From<TryCast> for PyTryCast {
64+
fn from(try_cast: TryCast) -> PyTryCast {
65+
PyTryCast { try_cast }
66+
}
67+
}
68+
69+
#[pymethods]
70+
impl PyTryCast {
71+
fn expr(&self) -> PyResult<PyExpr> {
72+
Ok((*self.try_cast.expr).clone().into())
73+
}
74+
75+
fn data_type(&self) -> PyResult<PyDataType> {
76+
Ok(self.try_cast.data_type.clone().into())
77+
}
78+
}

0 commit comments

Comments
 (0)
0