fix: Raise empty struct (pola-rs#17736)

atigbadr · Jul 23, 2024 · a1d519c · a1d519c
1 parent 656a453
commit a1d519c
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 4 deletions.
diff --git a/crates/polars-plan/src/dsl/functions/coerce.rs b/crates/polars-plan/src/dsl/functions/coerce.rs
@@ -1,7 +1,13 @@
 use super::*;
 
 /// Take several expressions and collect them into a [`StructChunked`].
+/// # Panics
+/// panics if `exprs` is empty.
 pub fn as_struct(exprs: Vec<Expr>) -> Expr {
+    assert!(
+        !exprs.is_empty(),
+        "expected at least 1 field in 'as_struct'"
+    );
     Expr::Function {
         input: exprs,
         function: FunctionExpr::AsStruct,

diff --git a/py-polars/polars/functions/as_datatype.py b/py-polars/polars/functions/as_datatype.py
@@ -596,15 +596,18 @@ def struct(
     Schema({'my_struct': Struct({'p': Int64, 'q': Boolean})})
     """
     pyexprs = parse_into_list_of_expressions(*exprs, **named_exprs)
-    expr = wrap_expr(plr.as_struct(pyexprs))
 
     if schema:
         if not exprs:
             # no columns or expressions provided; create one from schema keys
             expr = wrap_expr(
                 plr.as_struct(parse_into_list_of_expressions(list(schema.keys())))
             )
+        else:
+            expr = wrap_expr(plr.as_struct(pyexprs))
         expr = expr.cast(Struct(schema), strict=False)
+    else:
+        expr = wrap_expr(plr.as_struct(pyexprs))
 
     if eager:
         return F.select(expr).to_series()

diff --git a/py-polars/src/functions/lazy.rs b/py-polars/src/functions/lazy.rs
@@ -1,7 +1,7 @@
 use polars::lazy::dsl;
 use polars::prelude::*;
 use polars_plan::prelude::UnionArgs;
-use pyo3::exceptions::PyTypeError;
+use pyo3::exceptions::{PyTypeError, PyValueError};
 use pyo3::prelude::*;
 use pyo3::types::{PyBool, PyBytes, PyFloat, PyInt, PyString};
 
@@ -85,9 +85,14 @@ pub fn arg_where(condition: PyExpr) -> PyExpr {
 }
 
 #[pyfunction]
-pub fn as_struct(exprs: Vec<PyExpr>) -> PyExpr {
+pub fn as_struct(exprs: Vec<PyExpr>) -> PyResult<PyExpr> {
     let exprs = exprs.to_exprs();
-    dsl::as_struct(exprs).into()
+    if exprs.is_empty() {
+        return Err(PyValueError::new_err(
+            "expected at least 1 expression in 'as_struct'",
+        ));
+    }
+    Ok(dsl::as_struct(exprs).into())
 }
 
 #[pyfunction]

diff --git a/py-polars/tests/unit/datatypes/test_struct.py b/py-polars/tests/unit/datatypes/test_struct.py
@@ -965,3 +965,8 @@ def test_struct_out_nullability_from_arrow() -> None:
     df = pl.DataFrame(pd.DataFrame({"abc": [{"a": 1.0, "b": pd.NA}, pd.NA]}))
     res = df.select(a=pl.col("abc").struct.field("a"))
     assert res.to_dicts() == [{"a": 1.0}, {"a": None}]
+
+
+def test_empty_struct_raise() -> None:
+    with pytest.raises(ValueError):
+        pl.struct()