| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| from __future__ import annotations |
| |
| from datetime import datetime, timezone |
| |
| import pyarrow as pa |
| import pyarrow.compute as pc |
| import pytest |
| from datafusion import Accumulator, SessionContext, column, udaf |
| from datafusion.expr import Window |
| |
| |
| class Summarize(Accumulator): |
| """Interface of a user-defined accumulation.""" |
| |
| def __init__(self, initial_value: float = 0.0, as_scalar: bool = False): |
| self._sum = initial_value |
| self.as_scalar = as_scalar |
| |
| def state(self) -> list[pa.Scalar]: |
| if self.as_scalar: |
| return [pa.scalar(self._sum)] |
| return [self._sum] |
| |
| def update(self, values: pa.Array) -> None: |
| # Not nice since pyarrow scalars can't be summed yet. |
| # This breaks on `None` |
| self._sum = self._sum + pc.sum(values).as_py() |
| |
| def merge(self, states: list[pa.Array]) -> None: |
| # Not nice since pyarrow scalars can't be summed yet. |
| # This breaks on `None` |
| self._sum = self._sum + pc.sum(states[0]).as_py() |
| |
| def evaluate(self) -> pa.Scalar: |
| if self.as_scalar: |
| return pa.scalar(self._sum) |
| return self._sum |
| |
| |
| class NotSubclassOfAccumulator: |
| pass |
| |
| |
| class MissingMethods(Accumulator): |
| def __init__(self): |
| self._sum = pa.scalar(0) |
| |
| def state(self) -> list[pa.Scalar]: |
| return [self._sum] |
| |
| |
| class CollectTimestamps(Accumulator): |
| def __init__(self, wrap_in_scalar: bool): |
| self._values: list[datetime] = [] |
| self.wrap_in_scalar = wrap_in_scalar |
| |
| def state(self) -> list[pa.Scalar]: |
| if self.wrap_in_scalar: |
| return [pa.scalar(self._values, type=pa.list_(pa.timestamp("ns")))] |
| return [pa.array(self._values, type=pa.timestamp("ns"))] |
| |
| def update(self, values: pa.Array) -> None: |
| self._values.extend(values.to_pylist()) |
| |
| def merge(self, states: list[pa.Array]) -> None: |
| for state in states[0].to_pylist(): |
| if state is not None: |
| self._values.extend(state) |
| |
| def evaluate(self) -> pa.Scalar: |
| if self.wrap_in_scalar: |
| return pa.scalar(self._values, type=pa.list_(pa.timestamp("ns"))) |
| return pa.array(self._values, type=pa.timestamp("ns")) |
| |
| |
| @pytest.fixture |
| 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]], name="test_table") |
| |
| |
| def test_errors(df): |
| with pytest.raises(TypeError): |
| udaf( |
| NotSubclassOfAccumulator, |
| pa.float64(), |
| pa.float64(), |
| [pa.float64()], |
| volatility="immutable", |
| ) |
| |
| msg = ( |
| "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): |
| accum = udaf( # noqa: F841 |
| MissingMethods, |
| pa.int64(), |
| pa.int64(), |
| [pa.int64()], |
| volatility="immutable", |
| ) |
| |
| |
| def test_udaf_aggregate(df): |
| summarize = udaf( |
| Summarize, |
| pa.float64(), |
| pa.float64(), |
| [pa.float64()], |
| volatility="immutable", |
| ) |
| |
| df1 = df.aggregate([], [summarize(column("a"))]) |
| |
| # execute and collect the first (and only) batch |
| result = df1.collect()[0] |
| |
| assert result.column(0) == pa.array([1.0 + 2.0 + 3.0]) |
| |
| df2 = df.aggregate([], [summarize(column("a"))]) |
| |
| # Run a second time to ensure the state is properly reset |
| result = df2.collect()[0] |
| |
| assert result.column(0) == pa.array([1.0 + 2.0 + 3.0]) |
| |
| |
| def test_udaf_decorator_aggregate(df): |
| @udaf(pa.float64(), pa.float64(), [pa.float64()], "immutable") |
| def summarize(): |
| return Summarize() |
| |
| df1 = df.aggregate([], [summarize(column("a"))]) |
| |
| # execute and collect the first (and only) batch |
| result = df1.collect()[0] |
| |
| assert result.column(0) == pa.array([1.0 + 2.0 + 3.0]) |
| |
| df2 = df.aggregate([], [summarize(column("a"))]) |
| |
| # Run a second time to ensure the state is properly reset |
| result = df2.collect()[0] |
| |
| assert result.column(0) == pa.array([1.0 + 2.0 + 3.0]) |
| |
| |
| def test_udaf_decorator_keyword_arguments(df): |
| @udaf( |
| input_types=pa.float64(), |
| return_type=pa.float64(), |
| state_type=[pa.float64()], |
| volatility="immutable", |
| ) |
| def summarize(): |
| return Summarize() |
| |
| result = df.aggregate([], [summarize(column("a"))]).collect()[0] |
| assert result.column(0) == pa.array([1.0 + 2.0 + 3.0]) |
| |
| |
| def test_udaf_function_keyword_arguments(df): |
| summarize = udaf( |
| accum=Summarize, |
| input_types=pa.float64(), |
| return_type=pa.float64(), |
| state_type=[pa.float64()], |
| volatility="immutable", |
| ) |
| result = df.aggregate([], [summarize(column("a"))]).collect()[0] |
| assert result.column(0) == pa.array([1.0 + 2.0 + 3.0]) |
| |
| |
| @pytest.mark.parametrize("as_scalar", [True, False]) |
| def test_udaf_aggregate_with_arguments(df, as_scalar): |
| bias = 10.0 |
| |
| summarize = udaf( |
| lambda: Summarize(initial_value=bias, as_scalar=as_scalar), |
| pa.float64(), |
| pa.float64(), |
| [pa.float64()], |
| volatility="immutable", |
| ) |
| |
| df1 = df.aggregate([], [summarize(column("a"))]) |
| |
| # execute and collect the first (and only) batch |
| result = df1.collect()[0] |
| |
| assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0]) |
| |
| df2 = df.aggregate([], [summarize(column("a"))]) |
| |
| # Run a second time to ensure the state is properly reset |
| result = df2.collect()[0] |
| |
| assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0]) |
| |
| |
| def test_udaf_decorator_aggregate_with_arguments(df): |
| bias = 10.0 |
| |
| @udaf(pa.float64(), pa.float64(), [pa.float64()], "immutable") |
| def summarize(): |
| return Summarize(bias) |
| |
| df1 = df.aggregate([], [summarize(column("a"))]) |
| |
| # execute and collect the first (and only) batch |
| result = df1.collect()[0] |
| |
| assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0]) |
| |
| df2 = df.aggregate([], [summarize(column("a"))]) |
| |
| # Run a second time to ensure the state is properly reset |
| result = df2.collect()[0] |
| |
| assert result.column(0) == pa.array([bias + 1.0 + 2.0 + 3.0]) |
| |
| |
| def test_group_by(df): |
| summarize = udaf( |
| Summarize, |
| pa.float64(), |
| pa.float64(), |
| [pa.float64()], |
| volatility="immutable", |
| ) |
| |
| df = df.aggregate([column("b")], [summarize(column("a"))]) |
| |
| batches = df.collect() |
| |
| 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 |
| |
| |
| @pytest.fixture |
| def distinct_ctx(): |
| ctx = SessionContext() |
| ctx.register_udaf( |
| udaf( |
| Summarize, |
| pa.float64(), |
| pa.float64(), |
| [pa.float64()], |
| volatility="immutable", |
| ) |
| ) |
| ctx.from_pydict({"v": [1.0, 1.0, 1.0, 5.0]}, name="t") |
| return ctx |
| |
| |
| @pytest.mark.parametrize( |
| "run", |
| [ |
| pytest.param( |
| lambda ctx, summarize: ctx.table("t").aggregate( |
| [], [summarize(column("v")).distinct().build().alias("r")] |
| ), |
| id="distinct aggregate", |
| ), |
| pytest.param( |
| lambda ctx, summarize: ctx.table("t").select( |
| summarize(column("v")).over(Window()).distinct().build().alias("r") |
| ), |
| id="distinct window", |
| ), |
| ], |
| ) |
| def test_udaf_distinct_raises(distinct_ctx, run): |
| # The Python accumulator cannot deduplicate, so it would count every row. |
| summarize = udaf( |
| Summarize, |
| pa.float64(), |
| pa.float64(), |
| [pa.float64()], |
| volatility="immutable", |
| ) |
| with pytest.raises(Exception, match="DISTINCT is not supported"): |
| run(distinct_ctx, summarize).collect() |
| |
| |
| def test_udaf_distinct_rewritten_to_group_by_runs(distinct_ctx): |
| # The optimizer groups by the distinct values, so the accumulator never |
| # sees DISTINCT. |
| result = distinct_ctx.sql("select summarize(distinct v) as r from t") |
| assert result.collect_column("r").to_pylist() == [6.0] |
| |
| |
| @pytest.mark.parametrize("wrap_in_scalar", [True, False]) |
| def test_udaf_list_timestamp_return(ctx, wrap_in_scalar) -> None: |
| timestamps1 = [ |
| datetime(2024, 1, 1, tzinfo=timezone.utc), |
| datetime(2024, 1, 2, tzinfo=timezone.utc), |
| ] |
| timestamps2 = [ |
| datetime(2024, 1, 3, tzinfo=timezone.utc), |
| datetime(2024, 1, 4, tzinfo=timezone.utc), |
| ] |
| batch1 = pa.RecordBatch.from_arrays( |
| [pa.array(timestamps1, type=pa.timestamp("ns"))], |
| names=["ts"], |
| ) |
| batch2 = pa.RecordBatch.from_arrays( |
| [pa.array(timestamps2, type=pa.timestamp("ns"))], |
| names=["ts"], |
| ) |
| df = ctx.create_dataframe([[batch1], [batch2]], name="timestamp_table") |
| |
| list_type = pa.list_( |
| pa.field("item", type=pa.timestamp("ns"), nullable=wrap_in_scalar) |
| ) |
| |
| collect = udaf( |
| lambda: CollectTimestamps(wrap_in_scalar), |
| pa.timestamp("ns"), |
| list_type, |
| [list_type], |
| volatility="immutable", |
| ) |
| |
| result = df.aggregate([], [collect(column("ts"))]).collect()[0] |
| |
| # There is no guarantee about the ordering of the batches, so perform a sort |
| # to get consistent results. Alternatively we could sort on evaluate(). |
| assert ( |
| result.column(0).values.sort() |
| == pa.array( |
| [[*timestamps1, *timestamps2]], |
| type=list_type, |
| ).values |
| ) |
