How to implement a custom unnest function without getting "UDF returned a different number of rows than expected" error

[finlaydotb]

An unnest function by definition takes one value and has the possibility to generate multiple rows.

> select unnest(make_array(1, 2, 3, 4, 5)) as unnested;
+----------+
| unnested |
+----------+
| 1        |
| 2        |
| 3        |
| 4        |
| 5        |
+----------+

For some usecase I am trying to implement a custom unnest but I keep running into this error

called `Result::unwrap()` on an `Err` value: Internal("UDF returned a different number of rows than expected. Expected: 3, Got: 6")

The implementation looks something like this. This is not the custom logic I need but just to illustrate the issue

    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
        let args = ColumnarValue::values_to_arrays(args)?;
        let values = datafusion::common::cast::as_int32_array(&args[0])?;
        let mut builder = Int32Array::builder(args[0].len());

        values.into_iter().for_each(|value| {
            match value {
                Some(value) => builder.append_values(&vec![value, value], &vec![true, true]),
                None => builder.append_null(),
            }
        });

        Ok(ColumnarValue::Array(
            Arc::new(builder.finish()) as ArrayRef,
        ))
    }

[jonahgao]

Based on the definition of scalar functions:

SQL Scalar Functions operate on a single value and return a single value.

So args and results must have the same number of rows.

If it is not this case, it might be necessary to define it as a table function or custom logical plan. Unnest is a logical plan in DataFusion.

[finlaydotb]

If it is not this case, it might be necessary to define it as a table function or custom logical plan. Unnest is a logical plan in DataFusion.

Yes. This is what I concluded also.

necessary to define it as a table function

By this you mean this https://datafusion.apache.org/library-user-guide/adding-udfs.html#adding-a-user-defined-table-function ?
if so It does not look like UDTF can be used with values from columns in another table.

custom logical plan.

Do you by chance know any resources on how to do this? Not sure how a custom logical plan can be used to create a custom key work like my_unnest that can be used in a query.

[jonahgao]

if so It does not look like UDTF can be used with values from columns in another table.

Lateral joins can do this, but DataFusion does not support them yet.

Do you by chance know any resources on how to do this?

One possible way is to first declare my_unnest as a scalar UDF, then search for it in the logical plan, and transform the logical plan similarly to the approach in user_defined_plan.rs and the processing of UNNSET. But it certainly requires a lot of effort.