This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
The following commit(s) were added to refs/heads/main by this push:
new 044bbe2 Fix regression on register_udaf (#878)
044bbe2 is described below
commit 044bbe2f32a208f3e745971468d8769ca82e5d70
Author: Tim Saucer <[email protected]>
AuthorDate: Sat Sep 21 10:33:51 2024 -0400
Fix regression on register_udaf (#878)
* Test no longer hangs, and updated error string to match latest
* Add unit tests for registering udf and udaf
* Resolve error on registering udaf #874
* remove stale comment
* Update unit test text to match in multiple versions of python
* Regex for exception that is compatible with python 3.10 and 3.12
---
python/datafusion/tests/test_udaf.py | 47 ++++++++++++++++++++++++++++++------
python/datafusion/udf.py | 4 +--
2 files changed, 41 insertions(+), 10 deletions(-)
diff --git a/python/datafusion/tests/test_udaf.py
b/python/datafusion/tests/test_udaf.py
index 76488e1..6f2525b 100644
--- a/python/datafusion/tests/test_udaf.py
+++ b/python/datafusion/tests/test_udaf.py
@@ -21,7 +21,7 @@ import pyarrow as pa
import pyarrow.compute as pc
import pytest
-from datafusion import Accumulator, SessionContext, column, udaf
+from datafusion import Accumulator, column, udaf, udf
class Summarize(Accumulator):
@@ -60,18 +60,15 @@ class MissingMethods(Accumulator):
@pytest.fixture
-def df():
- ctx = SessionContext()
-
+def df(ctx):
# create a RecordBatch and a new DataFrame from it
batch = pa.RecordBatch.from_arrays(
[pa.array([1, 2, 3]), pa.array([4, 4, 6])],
names=["a", "b"],
)
- return ctx.create_dataframe([[batch]])
+ return ctx.create_dataframe([[batch]], name="test_table")
[email protected](reason="df.collect() will hang, need more investigations")
def test_errors(df):
with pytest.raises(TypeError):
udaf(
@@ -92,8 +89,9 @@ def test_errors(df):
df = df.aggregate([], [accum(column("a"))])
msg = (
- "Can't instantiate abstract class MissingMethods with abstract "
- "methods evaluate, merge, update"
+ "Can't instantiate abstract class MissingMethods (without an
implementation "
+ "for abstract methods 'evaluate', 'merge', 'update'|with abstract
methods "
+ "evaluate, merge, update)"
)
with pytest.raises(Exception, match=msg):
df.collect()
@@ -132,3 +130,36 @@ def test_group_by(df):
arrays = [batch.column(1) for batch in batches]
joined = pa.concat_arrays(arrays)
assert joined == pa.array([1.0 + 2.0, 3.0])
+
+
+def test_register_udaf(ctx, df) -> None:
+ summarize = udaf(
+ Summarize,
+ pa.float64(),
+ pa.float64(),
+ [pa.float64()],
+ volatility="immutable",
+ )
+
+ ctx.register_udaf(summarize)
+
+ df_result = ctx.sql("select summarize(b) from test_table")
+
+ assert df_result.collect()[0][0][0].as_py() == 14.0
+
+
+def test_register_udf(ctx, df) -> None:
+ is_null = udf(
+ lambda x: x.is_null(),
+ [pa.float64()],
+ pa.bool_(),
+ volatility="immutable",
+ name="is_null",
+ )
+
+ ctx.register_udf(is_null)
+
+ df_result = ctx.sql("select is_null(a) from test_table")
+ result = df_result.collect()[0].column(0)
+
+ assert result == pa.array([False, False, False])
diff --git a/python/datafusion/udf.py b/python/datafusion/udf.py
index a3b74bb..f74d675 100644
--- a/python/datafusion/udf.py
+++ b/python/datafusion/udf.py
@@ -192,7 +192,7 @@ class AggregateUDF:
See :py:func:`udaf` for a convenience function and argument
descriptions.
"""
- self._udf = df_internal.AggregateUDF(
+ self._udaf = df_internal.AggregateUDF(
name, accumulator, input_types, return_type, state_type,
str(volatility)
)
@@ -203,7 +203,7 @@ class AggregateUDF:
occur during the evaluation of the dataframe.
"""
args = [arg.expr for arg in args]
- return Expr(self._udf.__call__(*args))
+ return Expr(self._udaf.__call__(*args))
@staticmethod
def udaf(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
