Adding a pair of columns

[mikeprince4]

Hello -- I'm familiar with using @daft.udf to add a column to a DataFrame. However I need to add a pair of columns, the calculations of which depend on each other.

So something like this:

@daft.udf
def _calculate_column_pair(
    some_param: float,
    other_column_values: daft.series.Series,
) -> tuple[list[float], list[float]]:
    values_1 = []
    values_2 = []

    for other_col_val in other_column_values:
        value_1, value_2 = some_iterative_calculation(some_param, other_col_val)
        values_1.append(value_1)
        values_2.append(value_2)

    return values_1, values_2

I can't simply calculate the values_1 column and then calculate values_2 (or vice versa), due to the iterative nature of some_iterative_calculation which updates the values for value_1 and value_2 iteratively until a convergence criteria is met.

How best can I implement this in daft?

[mikeprince4]

update -- what I ended up doing was added a column as follows

@daft.udf(return_dtype=daft.DataType.fixed_size_list(daft.DataType.float64(), 2))

Then as a subsequent step, I split the column into two columns

df = df.with_column(
      first_element_column,
      daft.col(tuple_column).list.get(0),
  )

df = df.with_column(
      second_element_column,
      daft.col(tuple_column).list.get(1),
  )

[rchowell]

Yes, this or a struct is the current way for multi-column projections!

#4311

[rchowell]

Please feel free to re-open if there is additional help we can provide here 🙁