8000 feat: alias with metadata (#1111) · chenkovsky/datafusion-python@c9f1554 · GitHub
[go: up one dir, main page]
Skip to content

Commit c9f1554

Browse files
chenkovskychenkovsky
authored
feat: alias with metadata (apache#1111)
* feat: alias with metadata * fmt
1 parent d0d14f6 commit c9f1554

File tree

6 files changed

+44
-10
lines changed

6 files changed

+44
-10
lines changed

‎python/datafusion/expr.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -406,9 +406,17 @@ def column(value: str) -> Expr:
406406
"""Creates a new expression representing a column."""
407407
return Expr(expr_internal.RawExpr.column(value))
408408

409-
def alias(self, name: str) -> Expr:
410-
"""Assign a name to the expression."""
411-
return Expr(self.expr.alias(name))
409+
def alias(self, name: str, metadata: Optional[dict[str, str]] = None) -> Expr:
410+
"""Assign a name to the expression.
411+
412+
Args:
413+
name: The name to assign to the expression.
414+
metadata: Optional metadata to attach to the expression.
415+
416+
Returns:
417+
A new expression with the assigned name.
418+
"""
419+
return Expr(self.expr.alias(name, metadata))
412420

413421
def sort(self, ascending: bool = True, nulls_first: bool = True) -> SortExpr:
414422
"""Creates a sort :py:class:`Expr` from an existing :py:class:`Expr`.

‎python/datafusion/functions.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -372,9 +372,18 @@ def order_by(expr: Expr, ascending: bool = True, nulls_first: bool = True) -> So
372372
return SortExpr(expr, ascending=ascending, nulls_first=nulls_first)
373373

374374

375-
def alias(expr: Expr, name: str) -> Expr:
376-
"""Creates an alias expression."""
377-
return Expr(f.alias(expr.expr, name))
375+
def alias(expr: Expr, name: str, metadata: Optional[dict[str, str]] = None) -> Expr:
376+
"""Creates an alias expression with an optional metadata dictionary.
377+
378+
Args:
379+
expr: The expression to alias
380+
name: The alias name
381+
metadata: Optional metadata to attach to the column
382+
383+
Returns:
384+
An expression with the given alias
385+
"""
386+
return Expr(f.alias(expr.expr, name, metadata))
378387

379388

380389
def col(name: str) -> Expr:

‎python/tests/test_expr.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,3 +247,8 @@ def test_fill_null(df):
247247
assert result.column(0) == pa.array([1, 2, 100])
248248
assert result.column(1) == pa.array([4, 25, 6])
249249
assert result.column(2) == pa.array([1234, 1234, 8])
250+
251+
252+
def test_alias_with_metadata(df):
253+
df = df.select(col("a").alias("b", {"key": "value"}))
254+
assert df.schema().field("b").metadata == {b"key": b"value"}

‎python/tests/test_functions.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1231,3 +1231,8 @@ def test_between_default(df):
12311231

12321232
actual = df.collect()[0].to_pydict()
12331233
assert actual == expected
1234+
1235+
1236+
def test_alias_with_metadata(df):
1237+
df = df.select(f.alias(f.col("a"), "b", {"key": "value"}))
1238+
assert df.schema().field("b").metadata == {b"key": b"value"}

‎src/expr.rs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ use datafusion::logical_expr::{
2222
};
2323
use pyo3::IntoPyObjectExt;
2424
use pyo3::{basic::CompareOp, prelude::*};
25+
use std::collections::HashMap;
2526
use std::convert::{From, Into};
2627
use std::sync::Arc;
2728
use window::PyWindowFrame;
@@ -275,8 +276,9 @@ impl PyExpr {
275276
}
276277

277278
/// assign a name to the PyExpr
278-
pub fn alias(&self, name: &str) -> PyExpr {
279-
self.expr.clone().alias(name).into()
279+
#[pyo3(signature = (name, metadata=None))]
280+
pub fn alias(&self, name: &str, metadata: Option<HashMap<String, String>>) -> PyExpr {
281+
self.expr.clone().alias_with_metadata(name, metadata).into()
280282
}
281283

282284
/// Create a sort PyExpr from an existing PyExpr.

‎src/functions.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use std::collections::HashMap;
19+
1820
use datafusion::functions_aggregate::all_default_aggregate_functions;
1921
use datafusion::functions_window::all_default_window_functions;
2022
use datafusion::logical_expr::expr::WindowFunctionParams;
@@ -205,10 +207,13 @@ fn order_by(expr: PyExpr, asc: bool, nulls_first: bool) -> PyResult<PySortExpr>
205207

206208
/// Creates a new Alias Expr
207209
#[pyfunction]
208-
fn alias(expr: PyExpr, name: &str) -> PyResult<PyExpr> {
210+
#[pyo3(signature = (expr, name, metadata=None))]
211+
fn alias(expr: PyExpr, name: &str, metadata: Option<HashMap<String, String>>) -> PyResult<PyExpr> {
209212
let relation: Option<TableReference> = None;
210213
Ok(PyExpr {
211-
expr: datafusion::logical_expr::Expr::Alias(Alias::new(expr.expr, relation, name)),
214+
expr: datafusion::logical_expr::Expr::Alias(
215+
Alias::new(expr.expr, relation, name).with_metadata(metadata),
216+
),
212217
})
213218
}
214219

0 commit comments

Comments
 (0)
0