Skip to content

Commit a3c108f

Browse files
simicdsimicd
authored
docs: Example of calling Python UDF & UDAF in SQL (#258)
* Document UDF calls in SQL * Remove unnecessary imports * FIx example
1 parent 6cd74a9 commit a3c108f

File tree

4 files changed

+162
-2
lines changed

4 files changed

+162
-2
lines changed

‎README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,11 @@ See [examples](examples/README.md) for more information.
9595
- [Query a Parquet file using SQL](./examples/sql-parquet.py)
9696
- [Query a Parquet file using the DataFrame API](./examples/dataframe-parquet.py)
9797
- [Run a SQL query and store the results in a Pandas DataFrame](./examples/sql-to-pandas.py)
98+
- [Run a SQL query with a Python user-defined function (UDF)](./examples/sql-using-python-udf.py)
99+
- [Run a SQL query with a Python user-defined aggregation function (UDAF)](./examples/sql-using-python-udaf.py)
98100
- [Query PyArrow Data](./examples/query-pyarrow-data.py)
101+
- [Create dataframe](./examples/import.py)
102+
- [Export dataframe](./examples/export.py)
99103

100104
### Running User-Defined Python Code
101105

‎datafusion/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ def udf(func, input_types, return_type, volatility, name=None):
171171
if not callable(func):
172172
raise TypeError("`func` argument must be callable")
173173
if name is None:
174-
name = func.__qualname__
174+
name = func.__qualname__.lower()
175175
return ScalarUDF(
176176
name=name,
177177
func=func,
@@ -190,7 +190,7 @@ def udaf(accum, input_type, return_type, state_type, volatility, name=None):
190190
"`accum` must implement the abstract base class Accumulator"
191191
)
192192
if name is None:
193-
name = accum.__qualname__
193+
name = accum.__qualname__.lower()
194194
return AggregateUDF(
195195
name=name,
196196
accumulator=accum,

‎examples/sql-using-python-udaf.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
from datafusion import udaf, SessionContext, Accumulator
19+
import pyarrow as pa
20+
21+
22+
# Define a user-defined aggregation function (UDAF)
23+
class MyAccumulator(Accumulator):
24+
"""
25+
Interface of a user-defined accumulation.
26+
"""
27+
28+
def __init__(self):
29+
self._sum = pa.scalar(0.0)
30+
31+
def update(self, values: pa.Array) -> None:
32+
# not nice since pyarrow scalars can't be summed yet. This breaks on `None`
33+
self._sum = pa.scalar(
34+
self._sum.as_py() + pa.compute.sum(values).as_py()
35+
)
36+
37+
def merge(self, states: pa.Array) -> None:
38+
# not nice since pyarrow scalars can't be summed yet. This breaks on `None`
39+
self._sum = pa.scalar(
40+
self._sum.as_py() + pa.compute.sum(states).as_py()
41+
)
42+
43+
def state(self) -> pa.Array:
44+
return pa.array([self._sum.as_py()])
45+
46+
def evaluate(self) -> pa.Scalar:
47+
return self._sum
48+
49+
50+
my_udaf = udaf(
51+
MyAccumulator,
52+
pa.float64(),
53+
pa.float64(),
54+
[pa.float64()],
55+
"stable",
56+
# This will be the name of the UDAF in SQL
57+
# If not specified it will by default the same as accumulator class name
58+
name="my_accumulator",
59+
)
60+
61+
# Create a context
62+
ctx = SessionContext()
63+
64+
# Create a datafusion DataFrame from a Python dictionary
65+
source_df = ctx.from_pydict({"a": [1, 1, 3], "b": [4, 5, 6]})
66+
# Dataframe:
67+
# +---+---+
68+
# | a | b |
69+
# +---+---+
70+
# | 1 | 4 |
71+
# | 1 | 5 |
72+
# | 3 | 6 |
73+
# +---+---+
74+
75+
# Register UDF for use in SQL
76+
ctx.register_udaf(my_udaf)
77+
78+
# Query the DataFrame using SQL
79+
table_name = ctx.catalog().database().names().pop()
80+
result_df = ctx.sql(
81+
f"select a, my_accumulator(b) as b_aggregated from {table_name} group by a order by a"
82+
)
83+
# Dataframe:
84+
# +---+--------------+
85+
# | a | b_aggregated |
86+
# +---+--------------+
87+
# | 1 | 9 |
88+
# | 3 | 6 |
89+
# +---+--------------+
90+
assert result_df.to_pydict()["a"] == [1, 3]
91+
assert result_df.to_pydict()["b_aggregated"] == [9, 6]

‎examples/sql-using-python-udf.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
from datafusion import udf, SessionContext
19+
import pyarrow as pa
20+
21+
22+
# Define a user-defined function (UDF)
23+
def is_null(array: pa.Array) -> pa.Array:
24+
return array.is_null()
25+
26+
27+
is_null_arr = udf(
28+
is_null,
29+
[pa.int64()],
30+
pa.bool_(),
31+
"stable",
32+
# This will be the name of the UDF in SQL
33+
# If not specified it will by default the same as Python function name
34+
name="is_null",
35+
)
36+
37+
# Create a context
38+
ctx = SessionContext()
39+
40+
# Create a datafusion DataFrame from a Python dictionary
41+
source_df = ctx.from_pydict({"a": [1, 2, 3], "b": [4, None, 6]})
42+
# Dataframe:
43+
# +---+---+
44+
# | a | b |
45+
# +---+---+
46+
# | 1 | 4 |
47+
# | 2 | |
48+
# | 3 | 6 |
49+
# +---+---+
50+
51+
# Register UDF for use in SQL
52+
ctx.register_udf(is_null_arr)
53+
54+
# Query the DataFrame using SQL
55+
table_name = ctx.catalog().database().names().pop()
56+
result_df = ctx.sql(f"select a, is_null(b) as b_is_null from {table_name}")
57+
# Dataframe:
58+
# +---+-----------+
59+
# | a | b_is_null |
60+
# +---+-----------+
61+
# | 1 | false |
62+
# | 2 | true |
63+
# | 3 | false |
64+
# +---+-----------+
65+
assert result_df.to_pydict()["b_is_null"] == [False, True, False]

0 commit comments

Comments
 (0)