fix: Fix struct 'with_fields' schema for update dtypes (#16428)

pola-rs · May 23, 2024 · 30a5534 · 30a5534
1 parent f0e38b5
commit 30a5534
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 3 deletions.
diff --git a/crates/polars-plan/src/dsl/function_expr/struct_.rs b/crates/polars-plan/src/dsl/function_expr/struct_.rs
@@ -97,11 +97,20 @@ impl StructFunction {
                 let struct_ = &args[0];
 
                 if let DataType::Struct(fields) = struct_.data_type() {
-                    let mut fields = fields.iter().cloned().collect::<PlIndexSet<Field>>();
+                    let mut name_2_dtype = PlIndexMap::with_capacity(fields.len() * 2);
+
+                    for field in fields {
+                        name_2_dtype.insert(field.name(), field.data_type());
+                    }
                     for arg in &args[1..] {
-                        fields.insert(arg.clone());
+                        name_2_dtype.insert(arg.name(), arg.data_type());
                     }
-                    let dtype = DataType::Struct(fields.into_iter().collect());
+                    let dtype = DataType::Struct(
+                        name_2_dtype
+                            .iter()
+                            .map(|(name, dtype)| Field::new(name, (*dtype).clone()))
+                            .collect(),
+                    );
                     let mut out = struct_.clone();
                     out.coerce(dtype);
                     Ok(out)

diff --git a/py-polars/tests/unit/test_expansion.py b/py-polars/tests/unit/test_expansion.py
@@ -138,3 +138,15 @@ def test_struct_field_expand_rewrite() -> None:
     assert df.select(
         pl.struct(["A", "B"]).struct.field("*").name.prefix("foo_")
     ).to_dict(as_series=False) == {"foo_A": [1], "foo_B": [2]}
+
+
+def test_struct_field_expansion_16410() -> None:
+    q = pl.LazyFrame({"coords": [{"x": 4, "y": 4}]})
+
+    assert q.with_columns(
+        pl.col("coords").struct.with_fields(pl.field("x").sqrt()).struct.field("*")
+    ).collect().to_dict(as_series=False) == {
+        "coords": [{"x": 4, "y": 4}],
+        "x": [2.0],
+        "y": [4],
+    }