Open Wouittone opened 4 months ago
The key point is that the schema of FunctionExpr
in agg
seems doesn't has a list
type(in schema).
pl.LazyFrame({"a": [1,None], "b": [1,2]})
.group_by("b").agg(pl.col("a").drop_nulls())
.schema
OrderedDict([('b', Int64), ('a', Int64)])
@ritchie46 should have some insight.
I think function expr should call the to_field
with Context::Aggregation
. Not sure, but the schema is incorrect.
I think function expr should call the to_field with Context::Aggregation.
You mean when function get its input field? If we had agg(pl.col("a").foo())
, the input to foo()
would be a list
in Context::Aggregation
, which also feels problematic.
Maybe FunctionExpr
should handle Context
in function.get_field
(right now they are completely ignored). The only thing that bothering me is the return_scalar=True
case, which is somewhat propagate at the moment. For instance, agg(pl.col("a").sum().pow(2)))
is not nested even though pow
does not return_scalar
itself, so we can't simply wrap a list on its schema.
I was about to start working on this a bit further, but I now realize that I probably shouldn't have used a concat
in the aggregation step.
I just tried running the example again without it and it works just as expected:
(
pl.LazyFrame({"id": [1, 1, 2, 2, 2, 3, 3, 3, 3], "values": [1, 2, 3, 4, 5, 6, 7, 8, 9]})
.group_by("id")
.agg("values")
.select(pl.col("values").list.gather([0, 1]).list.to_array(2).alias("lists_start"))
.collect()
)
>>> shape: (3, 1)
┌───────────────┐
│ lists_start │
│ --- │
│ array[i64, 2] │
╞═══════════════╡
│ [6, 7] │
│ [1, 2] │
│ [3, 4] │
└───────────────┘
I personally feel that it would be fine to leave it as is, and maybe improve the documentation on this point if you think this is relevant? I'd be happy to give it a shot, if so.
(Sorry for the misclick... :disappointed:)
Checks
Reproducible example
Log output
Issue description
This is a strange error where applying
.list.to_array
in a LazyFrame after a grouping operation that concatenates some values into a list seems to wrongly infer the column type (at least that is what I understand from theComputeError: expected List dtype
error message.)Running the example without the array cast works and the column dtype is a
pl.List
:Expected behavior
This seems to work fine for eager DataFrame:
And also for a LazyFrame without a groupby:
I was expecting the LazyFrame with a grouby operation to work just as those two.
Installed versions