def test_udf(con_for_udf, test_schema, table):
"""Test creating a UDF in database based on Python function
and then creating an ibis UDF object based on that"""
mult_a_b_udf = udf(
con_for_udf,
mult_a_b,
(dt.int32, dt.int32),
dt.int32,
schema=test_schema,
replace=True,
)
table_filt = table.filter(table['user_id'] == 2)
expr = table_filt[mult_a_b_udf(
table_filt['user_id'], table_filt['name_length']).name('mult_result')]
result = expr.execute()
assert result['mult_result'].iloc[0] == 8
def test_array_type(con_for_udf, test_schema, table):
"""Test that usage of Array types work
Other scalar types can be represented either by the class or an instance,
but Array types work differently. Array types must be an instance,
because the Array class must be instantiated specifying the datatype
of the elements of the array.
"""
pysplit_udf = udf(
con_for_udf,
pysplit,
(dt.string, dt.string),
dt.Array(dt.string),
schema=test_schema,
replace=True,
)
splitter = ibis.literal(' ', dt.string)
result = pysplit_udf(table['user_name'], splitter).name('split_name')
result.execute()
def udf(self,
pyfunc,
in_types,
out_type,
schema=None,
replace=False,
name=None):
"""Decorator that defines a PL/Python UDF in-database based on the
wrapped function and turns it into an ibis function expression.
Parameters
----------
pyfunc : function
in_types : List[ibis.expr.datatypes.DataType]
out_type : ibis.expr.datatypes.DataType
schema : str
optionally specify the schema in which to define the UDF
replace : bool
replace UDF in database if already exists
name: str
name for the UDF to be defined in database
Returns
-------
Callable
Function that takes in ColumnExpr arguments and returns an instance
inheriting from PostgresUDFNode
"""
return udf(
client=self,
python_func=pyfunc,
in_types=in_types,
out_type=out_type,
schema=schema,
replace=replace,
name=name,
)