Open Baukebrenninkmeijer opened 1 month ago
This appears to be a minimal repro:
df = pl.DataFrame({"text_column": [None] * 2})
df.with_columns(
pl.all().map_elements(lambda x: {"id": None}, return_dtype=pl.Struct)
.cast(pl.Struct({"id": pl.Int64}))
)
# ShapeError: unable to add a column of length 0 to a DataFrame of height 2
There seems to be some problem by allowing an empty pl.Struct
in return_dtype
e.g. https://github.com/pola-rs/polars/issues/8141 https://github.com/pola-rs/polars/issues/17181
Just on the topic of the actual functionality, perhaps it could be done without map_elements
:
df.with_columns(
pl.struct(
pl.col('text_column')
.str.extract(r'\{([^}]+)\}')
.str.extract(rf'{name}=([^,}}]+)')
.cast(dtype, strict=False)
.alias(name)
for name, dtype in schema.items()
)
)
@cmdlineluser I'll have a look at this. Before I was iterating through each columns in the with_columns which worked for eager mode, but in lazy mode, all columns got the value of the latest column (possibly another bug). I'll see whether your example works, it would be much cleaner at least.
I am having the same issue. By specifying the fields, I managed to get it running.
from my_module import parse_parameter
import polars as pl
df = pl.DataFrame(
{'parameter1': ['[MS, MS:1001477, SpectraST,]',
'[MOD, MOD:00648, "N,O-diacetylated L-serine",]',
None],
'parameter2': [None, None, None]}
)
df.with_columns(
pl.all().map_elements(
parse_parameter,
return_dtype=pl.Struct(
fields=[pl.Field(name, pl.String) for name in 'abcd']
),
)
)
For the parameter1
column it works without fields, for paramter2
column not.
Is it possible somehow to have Nulls in Struct columns? This would be really useful to me since I am validating the Struct against a Pydantic model which can either be Null or should have all fields.
A question I do not know the answer to is:
What is the correct way to return a schema'd dict from a UDF?
def return_dict(s):
result = {"a": "1", "b": "two"}
return result
schema = {"a": pl.Int64, "b": pl.String}
df = pl.DataFrame({"x": 1})
>>> df.with_columns(y = pl.all().map_elements(return_dict, return_dtype=pl.Struct(schema)))
SchemaError: expected output type 'Struct([Field { name: "a", dtype: Int64 }, Field { name: "b", dtype: String }])', got 'Struct([Field { name: "a", dtype: String }, Field { name: "b", dtype: String }])'; set `return_dtype` to the proper datatype
Setting return_dtype=pl.Struct
allows you to bypass this error which you then cast: (should this be allowed?)
>>> df.with_columns(y = pl.all().map_elements(return_dict, return_dtype=pl.Struct).cast(pl.Struct(schema)))
shape: (1, 2)
┌─────┬───────────┐
│ x ┆ y │
│ --- ┆ --- │
│ i64 ┆ struct[2] │
╞═════╪═══════════╡
│ 1 ┆ {1,"two"} │
└─────┴───────────┘
But if there are no non-null values - you hit the ShapeError bug.
If we can modify the UDF, we can get a typed dict from a DataFrame, and set the correct return_dtype
:
def return_dict(s):
result = {"a": "1", "b": "two"}
return pl.DataFrame(result).cast(schema).to_struct().item()
schema = {"a": pl.Int64, "b": pl.String}
df = pl.DataFrame({"x": 1})
df.with_columns(y = pl.all().map_elements(return_dict, return_dtype=pl.Struct(schema)))
# shape: (1, 2)
# ┌─────┬───────────┐
# │ x ┆ y │
# │ --- ┆ --- │
# │ i64 ┆ struct[2] │
# ╞═════╪═══════════╡
# │ 1 ┆ {1,"two"} │
# └─────┴───────────┘
(pl.DataFrame({"x": [None]})
.with_columns(y = pl.all().map_elements(return_dict, return_dtype=pl.Struct(schema)))
)
# shape: (1, 2)
# ┌──────┬─────────────┐
# │ x ┆ y │
# │ --- ┆ --- │
# │ null ┆ struct[2] │
# ╞══════╪═════════════╡
# │ null ┆ {null,null} │
# └──────┴─────────────┘
But I'm not sure if this is the way this is supposed to be done?
I would be very surprised if this was the intended way. But it's good to have this as a workaround.
@janikkokot So putting the schema in fields rather than a dict works to circumvent this problem?
@cmdlineluser
df.with_columns(
pl.struct(
pl.col('text_column')
.str.extract(r'\{([^}]+)\}')
.str.extract(rf'{name}=([^,}}]+)')
.cast(dtype, strict=False)
.alias(name)
for name, dtype in schema.items()
)
)
This works flawlessly, thanks for suggesting it!
Checks
Reproducible example
Log output
Issue description
When using this structure in a
with_columns
col, the length misalignment causes this to fail.Expected behavior
My expectation is, even with
skip_nulls=True
, would be that the first one still returns a Series with the shape (4, 1), filled with Nulls.Installed versions