docs: Not clear how to define udfs on struct types

[vspinu]

Please describe the issue

I cannot figure out a way to define udfs with struct input.

This is my attempt with spark:

df = pd.DataFrame({"a": [1, 2, 3]})
sdf = spark.createDataFrame(df)
sdf.createOrReplaceTempView("df")

con = ibis.pyspark.connect(spark)
tbl = con.table("df")

SelectorType = ibis.expr.datatypes.Struct({"a": str})

@ibis.udf.scalar.pandas
def udf(x: SelectorType) -> int:
     return 1

tbl.mutate(struct=ibis.struct({"a": _.a})).mutate(struct2=udf(_.struct)).execute()

I get the following signature error:

SignatureValidationError: udf_3(r0 := DatabaseTable: df
  a int64

r1 := Project[r0]
  a:      r0.a
  struct: StructColumn(names=['a'], values=[r0.a])

struct: r1.struct) has failed due to the following errors:
  `x`: r0 := DatabaseTable: df
  a int64

r1 := Project[r0]
  a:      r0.a
  struct: StructColumn(names=['a'], values=[r0.a])

struct: r1.struct of type <class 'ibis.expr.types.structs.StructColumn'> is not matching ValueOf(dtype='Struct([('a', String(length=None, nullable=True))], nullable=True))

Expected signature: udf_3(x: struct<a: string>)

Would it be possible to provide an example in the official docs of how to build a custom UDF with struct input and/or output types?

Thanks!

Code of Conduct

• I agree to follow this project's Code of Conduct

[cpcloud]

Your UDF is defined as accepting a struct with string field:

SelectorType = ibis.expr.datatypes.Struct({"a": str})

@ibis.udf.scalar.pandas
def udf(x: SelectorType) -> int:
     return 1

but you're later constructing a struct with an integer field, because you're using the data from df.a:

df = pd.DataFrame({"a": [1, 2, 3]})
# ...
tbl.mutate(struct=ibis.struct({"a": _.a}))

You can either redefined SelectorType, or you can cast _.a to a string with _.a.cast("str") to match the expected type you declared in your UDF.

---

You also need to return a Series from your UDF (the int is the element type not the literal Python return type):

@ibis.udf.scalar.pandas
def udf(x: SelectorType) -> int:
     return pd.Series([1] * len(x))

This last thing can be documented better.

---

Once you make those changes, this works:

In [34]: from ibis.interactive import *

In [35]: import pandas as pd

In [36]: from pyspark.sql import SparkSession

In [37]: spark = SparkSession.builder.getOrCreate()

In [38]: df = pd.DataFrame({"a": [1, 2, 3]})
    ...: sdf = spark.createDataFrame(df)
    ...: sdf.createOrReplaceTempView("df")
    ...:
    ...: con = ibis.pyspark.connect(spark)
    ...: tbl = con.table("df")
    ...:
    ...: SelectorType = ibis.expr.datatypes.Struct({"a": str})
    ...:
    ...: @ibis.udf.scalar.pandas
    ...: def udf(x: SelectorType) -> int:
    ...:      return pd.Series([1] * len(x))
    ...:

In [39]: tbl.mutate(struct=ibis.struct({"a": _.a.cast("str")})).mutate(struct2=udf(_.struct)).execute()
Out[39]:
   a      struct  struct2
0  1  {'a': '1'}        1
1  2  {'a': '2'}        1
2  3  {'a': '3'}        1