8000 add PyUnnestExpr wrapper · apache/datafusion-python@ffed35d · GitHub
[go: up one dir, main page]

Skip to content

Commit ffed35d

Browse files
add PyUnnestExpr wrapper
1 parent 0d571ab commit ffed35d

File tree

2 files changed

+70
-1
lines changed

2 files changed

+70
-1
lines changed

src/expr.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ pub mod subquery_alias;
9090
pub mod table_scan;
9191
pub mod union;
9292
pub mod unnest;
93+
pub mod unnest_expr;
9394
pub mod window;
9495

9596
/// A PyExpr that can be used on a DataFrame
@@ -169,7 +170,7 @@ impl PyExpr {
169170
Ok(placeholder::PyPlaceholder::from(value.clone()).into_py(py))
170171
}
171172
Expr::OuterReferenceColumn(_, _) => todo!(),
172-
Expr::Unnest(_) => todo!(),
173+
Expr::Unnest(value) => Ok(unnest_expr::PyUnnestExpr::from(value.clone()).into_py(py)),
173174
})
174175
}
175176

@@ -624,6 +625,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
624625
m.add_class::<cross_join::PyCrossJoin>()?;
625626
m.add_class::<union::PyUnion>()?;
626627
m.add_class::<unnest::PyUnnest>()?;
628+
m.add_class::<unnest_expr::PyUnnestExpr>()?;
627629
m.add_class::<extension::PyExtension>()?;
628630
m.add_class::<filter::PyFilter>()?;
629631
m.add_class::<projection::PyProjection>()?;

src/expr/unnest_expr.rs

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
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 datafusion_expr::expr::Unnest;
19+
use pyo3::prelude::*;
20+
use std::fmt::{self, Display, Formatter};
21+
22+
use super::PyExpr;
23+
24+
#[pyclass(name = "UnnestExpr", module = "datafusion.expr", subclass)]
25+
#[derive(Clone)]
26+
pub struct PyUnnestExpr {
27+
unnest: Unnest,
28+
}
29+
30+
impl From<Unnest> for PyUnnestExpr {
31+
fn from(unnest: Unnest) -> PyUnnestExpr {
32+
PyUnnestExpr { unnest }
33+
}
34+
}
35+
36+
impl From<PyUnnestExpr> for Unnest {
37+
fn from(unnest: PyUnnestExpr) -> Self {
38+
unnest.unnest
39+
}
40+
}
41+
42+
impl Display for PyUnnestExpr {
43+
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
44+
write!(
45+
f,
46+
"Unnest
47+
Expr: {:?}",
48+
&self.unnest.expr,
49+
)
50+
}
51+
}
52+
53+
#[pymethods]
54+
impl PyUnnestExpr {
55+
/// Retrieves the expression that is being unnested
56+
fn expr(&self) -> PyResult<PyExpr> {
57+
Ok((*self.unnest.expr).clone().into())
58+
}
59+
60+
fn __repr__(&self) -> PyResult<String> {
61+
Ok(format!("UnnestExpr({})", self))
62+
}
63+
64+
fn __name__(&self) -> PyResult<String> {
65+
Ok("UnnestExpr".to_string())
66+
}
67+
}

0 commit comments

Comments
 (0)
0