Is it possible to perform multiple aggregations ov...
# general
t
Is it possible to perform multiple aggregations over a group? The udfs I'm trying to apply seem to work if I supply one of them to
.agg
after a
groupby
, however, if I supply multiple expressions to
.agg
I get the error:
Copy code
DaftError::ValueError MapGroups not supported via aggregation, use map_groups instead
But
map_groups
looks like it only accepts 1 expression.
k
Hey Tyler, what you can do is return a StructType in your
map_groups
function with the aggregation columns that you want, and then expand that struct out after the aggregation with
df.select("struct_col.*")
t
Got it. That works.
This seems to work for everything but tensor types, is that a known issue?
k
It should work. What's the error that you see?
t
NotImplemented
from rust
k
Ah. Does it tell you the line in Rust?
t
Let me try running it with traces
This one is a particularly hard one to capture because it seems like ray treats it as retryable so it generates a lot of stacktraces.
Sorry for the length of this one but I wasn't sure which parts might be relevant.
The udf in question looks like this:
Copy code
@daft.udf(
    return_dtype=DT.struct(
        {
            "positions": DT.tensor(DT.float64()),
            "expressions": DT.tensor(DT.float64()),
            "masks": DT.list(DT.bool()),
        }
    )
)
def format_experiment_inputs(
    x_diff: daft.Series,
    y_diff: daft.Series,
    expression: daft.Series,
    length: int = 6,
):
    return [
        {
            "positions": build_positions(x_diff, y_diff, length),
            "expressions": build_expressions(expression, length),
            "masks": build_missing_data_mask(x_diff, length),
        }
    ]
If I make all of those smaller functions individual udfs with the same inner type then it works.
But then I have to fan out the grouping to multiple
map_group
calls and then join the outputs which is more overhead.
j
I think it’s here (casting of Python arrays to Structs) We should ban any usage of
unimplemented!()
without a useful message 😛 cc @Desmond Cheong as well
plus one 2
Ok managed to reproduce:
Copy code
import numpy as np
import daft
from daft import DataType as DT

@daft.udf(
    return_dtype=DT.struct(
        {
            "positions": DT.tensor(DT.float64()),
            "expressions": DT.tensor(DT.float64()),
            "masks": DT.list(DT.bool()),
        }
    )
)
def format_experiment_inputs(
    x_diff: daft.Series,
    y_diff: daft.Series,
    expression: daft.Series,
    length: int = 6,
):
    return [
        {
            "positions": np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
            "expressions": np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
            "masks": [[True, True, False], None, None],
        }
    ]

df = daft.from_pydict({"foo": [b"data1", b"data2", b"data3"]})
df = df.with_column("bar", format_experiment_inputs(df["foo"], df["foo"], df["foo"]))
df.collect()
I think our
numpy -> Series
story needs to be improved overall… Noticing some weird issues even without structs. Sorry about that @Tyler Van Hensbergen we’ll straighten this out and get back to you shortly @Desmond Cheong let’s sync up and come up with a good tracking issue for this so we can knock it out?
🫡 1
t
Awesome, and no need to apologize. Just want to emphasize that this one is not a critical blocker. I have a decent workaround, this would just speed up the performance a bit.
👍 1