blob: 3763c7cb88a3bb7268163926169d52db14644b63 [file]
# 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.
import numpy as np
import pyarrow as pa
import pytest
from datafusion import SessionContext, column, lit
from datafusion import functions as f
from datafusion.common import NullTreatment
@pytest.fixture
def df():
ctx = SessionContext()
# create a RecordBatch and a new DataFrame from it
batch = pa.RecordBatch.from_arrays(
[
pa.array([1, 2, 3]),
pa.array([4, 4, 6]),
pa.array([9, 8, 5]),
pa.array([True, True, False]),
pa.array([1, 2, None]),
],
names=["a", "b", "c", "d", "e"],
)
return ctx.create_dataframe([[batch]])
@pytest.fixture
def df_partitioned():
ctx = SessionContext()
# create a RecordBatch and a new DataFrame from it
batch = pa.RecordBatch.from_arrays(
[
pa.array([0, 1, 2, 3, 4, 5, 6]),
pa.array([7, None, 7, 8, 9, None, 9]),
pa.array(["A", "A", "A", "A", "B", "B", "B"]),
],
names=["a", "b", "c"],
)
return ctx.create_dataframe([[batch]])
@pytest.fixture
def df_aggregate_100():
ctx = SessionContext()
ctx.register_csv("aggregate_test_data", "./testing/data/csv/aggregate_test_100.csv")
return ctx.table("aggregate_test_data")
@pytest.mark.parametrize(
("agg_expr", "calc_expected"),
[
(f.avg(column("a")), lambda a, b, c, d: np.array(np.average(a))),
(
f.corr(column("a"), column("b")),
lambda a, b, c, d: np.array(np.corrcoef(a, b)[0][1]),
),
(f.count(column("a")), lambda a, b, c, d: pa.array([len(a)])),
# Sample (co)variance -> ddof=1
# Population (co)variance -> ddof=0
(
f.covar(column("a"), column("b")),
lambda a, b, c, d: np.array(np.cov(a, b, ddof=1)[0][1]),
),
(
f.covar_pop(column("a"), column("c")),
lambda a, b, c, d: np.array(np.cov(a, c, ddof=0)[0][1]),
),
(
f.covar_samp(column("b"), column("c")),
lambda a, b, c, d: np.array(np.cov(b, c, ddof=1)[0][1]),
),
# f.grouping(col_a), # noqa: ERA001 No physical plan implemented yet
(f.max(column("a")), lambda a, b, c, d: np.array(np.max(a))),
(f.mean(column("b")), lambda a, b, c, d: np.array(np.mean(b))),
(f.median(column("b")), lambda a, b, c, d: np.array(np.median(b))),
(f.min(column("a")), lambda a, b, c, d: np.array(np.min(a))),
(f.sum(column("b")), lambda a, b, c, d: np.array(np.sum(b.to_pylist()))),
# Sample stdev -> ddof=1
# Population stdev -> ddof=0
(f.stddev(column("a")), lambda a, b, c, d: np.array(np.std(a, ddof=1))),
(f.stddev_pop(column("b")), lambda a, b, c, d: np.array(np.std(b, ddof=0))),
(f.stddev_samp(column("c")), lambda a, b, c, d: np.array(np.std(c, ddof=1))),
(f.var(column("a")), lambda a, b, c, d: np.array(np.var(a, ddof=1))),
(f.var_pop(column("b")), lambda a, b, c, d: np.array(np.var(b, ddof=0))),
(f.var_samp(column("c")), lambda a, b, c, d: np.array(np.var(c, ddof=1))),
],
)
def test_aggregation_stats(df, agg_expr, calc_expected):
df = df.select("a", "b", "c", "d")
agg_df = df.aggregate([], [agg_expr])
result = agg_df.collect()[0]
values_a, values_b, values_c, values_d = df.collect()[0]
expected = calc_expected(values_a, values_b, values_c, values_d)
np.testing.assert_array_almost_equal(result.column(0), expected)
@pytest.mark.parametrize(
("agg_expr", "expected", "array_sort"),
[
(f.approx_distinct(column("b")), pa.array([2], type=pa.uint64()), False),
(
f.approx_distinct(
column("b"),
filter=column("a") != lit(3),
),
pa.array([1], type=pa.uint64()),
False,
),
(f.approx_median(column("b")), pa.array([4], type=pa.float64()), False),
(
f.median(column("b"), distinct=True),
pa.array([5], type=pa.float64()),
False,
),
(
f.median(column("b"), filter=column("a") != 2),
pa.array([5], type=pa.float64()),
False,
),
(
f.approx_median(column("b"), filter=column("a") != 2),
pa.array([5], type=pa.float64()),
False,
),
(
f.approx_percentile_cont(column("b"), 0.5),
pa.array([4], type=pa.float64()),
False,
),
(
f.approx_percentile_cont(
column("b").sort(ascending=True, nulls_first=False),
0.5,
num_centroids=2,
),
pa.array([4.75], type=pa.float64()),
False,
),
(
f.approx_percentile_cont_with_weight(column("b"), lit(0.6), 0.5),
pa.array([4], type=pa.float64()),
False,
),
(
f.approx_percentile_cont_with_weight(
column("b").sort(ascending=False, nulls_first=False), lit(0.6), 0.5
),
pa.array([4], type=pa.float64()),
False,
),
(
f.approx_percentile_cont_with_weight(
column("b"), lit(0.6), 0.5, filter=column("a") != lit(3)
),
pa.array([4], type=pa.float64()),
False,
),
(f.array_agg(column("b")), pa.array([[4, 4, 6]]), False),
(f.array_agg(column("b"), distinct=True), pa.array([[4, 6]]), True),
(
f.array_agg(column("e"), filter=column("e").is_not_null()),
pa.array([[1, 2]]),
False,
),
(
f.array_agg(column("b"), order_by=[column("c")]),
pa.array([[6, 4, 4]]),
False,
),
(
f.array_agg(column("b"), order_by=column("c")),
pa.array([[6, 4, 4]]),
False,
),
(f.avg(column("b"), filter=column("a") != lit(1)), pa.array([5.0]), False),
(f.avg(column("b"), distinct=True), pa.array([5.0]), False),
(f.sum(column("b"), filter=column("a") != lit(1)), pa.array([10]), False),
(f.sum(column("b"), distinct=True), pa.array([10]), False),
(f.count(column("b"), distinct=True), pa.array([2]), False),
(f.count(column("b"), filter=column("a") != 3), pa.array([2]), False),
(f.count(), pa.array([3]), False),
(f.count(column("e")), pa.array([2]), False),
(f.count_star(filter=column("a") != 3), pa.array([2]), False),
(f.max(column("a"), filter=column("a") != lit(3)), pa.array([2]), False),
(f.min(column("a"), filter=column("a") != lit(1)), pa.array([2]), False),
(
f.stddev(column("a"), filter=column("a") != lit(2)),
pa.array([np.sqrt(2)]),
False,
),
(
f.stddev_pop(column("a"), filter=column("a") != lit(2)),
pa.array([1.0]),
False,
),
],
)
def test_aggregation(df, agg_expr, expected, array_sort):
agg_df = df.aggregate([], [agg_expr.alias("agg_expr")])
if array_sort:
agg_df = agg_df.select(f.array_sort(column("agg_expr")))
agg_df.show()
result = agg_df.collect()[0]
assert result.column(0) == expected
@pytest.mark.parametrize(
("name", "expr", "expected"),
[
(
"approx_percentile_cont",
f.approx_percentile_cont(column("c3"), 0.95, num_centroids=200),
[73.55, 68.0, 122.5, 124.2, 115.6],
),
(
"approx_perc_cont_few_centroids",
f.approx_percentile_cont(column("c3"), 0.95, num_centroids=5),
[72.775, 68.0, 119.4075, 124.825, 115.44],
),
(
"approx_perc_cont_filtered",
f.approx_percentile_cont(
column("c3"), 0.95, num_centroids=200, filter=column("c3") > lit(0)
),
[83.0, 68.0, 122.75, 124.9, 117.6],
),
(
"corr",
f.corr(column("c3"), column("c2")),
[-0.1056, -0.2808, 0.0023, 0.0022, -0.2473],
),
(
"corr_w_filter",
f.corr(column("c3"), column("c2"), filter=column("c3") > lit(0)),
[-0.3298, 0.2925, 0.2467, -0.2269, 0.0358],
),
(
"covar_pop",
f.covar_pop(column("c3"), column("c2")),
[-7.2857, -25.6731, 0.2222, 0.2469, -20.2857],
),
(
"covar_pop_w_filter",
f.covar_pop(column("c3"), column("c2"), filter=column("c3") > lit(0)),
[-9.25, 9.0579, 13.7521, -9.9669, 1.1641],
),
(
"covar_samp",
f.covar_samp(column("c3"), column("c2")),
[-7.65, -27.0994, 0.2333, 0.2614, -21.3],
),
(
"covar_samp_w_filter",
f.covar_samp(column("c3"), column("c2"), filter=column("c3") > lit(0)),
[-10.5714, 9.9636, 15.1273, -10.9636, 1.2417],
),
(
"var_samp",
f.var_samp(column("c2")),
[1.9286, 2.2047, 1.6333, 2.1438, 1.6],
),
(
"var_samp_w_filter",
f.var_samp(column("c2"), filter=column("c3") > lit(0)),
[1.4286, 2.4182, 1.8545, 1.4727, 1.6292],
),
(
"var_pop",
f.var_pop(column("c2")),
[1.8367, 2.0886, 1.5556, 2.0247, 1.5238],
),
(
"var_pop_w_filter",
f.var_pop(column("c2"), filter=column("c3") > lit(0)),
[1.25, 2.1983, 1.686, 1.3388, 1.5273],
),
],
)
def test_aggregate_100(df_aggregate_100, name, expr, expected):
# https://github.com/apache/datafusion/blob/bddb6415a50746d2803dd908d19c3758952d74f9/datafusion/sqllogictest/test_files/aggregate.slt#L1490-L1498
df = (
df_aggregate_100.aggregate(
[column("c1")],
[expr.alias(name)],
)
.select("c1", f.round(column(name), lit(4)).alias(name))
.sort(column("c1").sort(ascending=True))
)
df.show()
expected_dict = {
"c1": ["a", "b", "c", "d", "e"],
name: expected,
}
assert df.collect()[0].to_pydict() == expected_dict
def test_any_value_skips_nulls_per_group():
ctx = SessionContext()
df = ctx.from_pydict(
{"g": ["x", "x", "y", "y", "z"], "v": [None, 7, 8, None, None]}
)
result = (
df.aggregate([column("g")], [f.any_value(column("v")).alias("v")])
.sort(column("g").sort())
.to_pydict()
)
assert result == {"g": ["x", "y", "z"], "v": [7, 8, None]}
@pytest.mark.parametrize(
("expr", "expected"),
[
pytest.param(f.mean(column("v")), 2.0, id="mean"),
pytest.param(f.mean(column("v"), distinct=True), 3.0, id="mean_distinct"),
pytest.param(
f.mean(column("v"), filter=column("v") > lit(1.0)), 5.0, id="mean_filter"
),
pytest.param(f.percentile_cont(column("v"), 0.5), 1.0, id="percentile_cont"),
pytest.param(
f.percentile_cont(column("v"), 0.5, distinct=True),
3.0,
id="percentile_cont_distinct",
),
pytest.param(
f.quantile_cont(column("v"), 0.5, distinct=True),
3.0,
id="quantile_cont_distinct",
),
],
)
def test_distinct_numeric_aggregates(expr, expected):
ctx = SessionContext()
df = ctx.from_pydict({"v": [1.0, 1.0, 1.0, 5.0]})
result = df.aggregate([], [expr.alias("r")]).collect_column("r")[0].as_py()
assert result == expected
data_test_bitwise_and_boolean_functions = [
("any_value_filter", f.any_value(column("a"), filter=column("a") == lit(2)), [2]),
("bit_and", f.bit_and(column("a")), [0]),
("bit_and_filter", f.bit_and(column("a"), filter=column("a") != lit(2)), [1]),
("bit_or", f.bit_or(column("b")), [6]),
("bit_or_filter", f.bit_or(column("b"), filter=column("a") != lit(3)), [4]),
("bit_and_distinct", f.bit_and(column("b"), distinct=True), [4]),
("bit_or_distinct", f.bit_or(column("b"), distinct=True), [6]),
("bit_xor", f.bit_xor(column("c")), [4]),
("bit_xor_distinct", f.bit_xor(column("b"), distinct=True), [2]),
("bit_xor_filter", f.bit_xor(column("b"), filter=column("a") != lit(3)), [0]),
(
"bit_xor_filter_distinct",
f.bit_xor(column("b"), distinct=True, filter=column("a") != lit(3)),
[4],
),
("bool_and", f.bool_and(column("d")), [False]),
("bool_and_filter", f.bool_and(column("d"), filter=column("a") != lit(3)), [True]),
("bool_or", f.bool_or(column("d")), [True]),
("bool_or_filter", f.bool_or(column("d"), filter=column("a") == lit(3)), [False]),
]
@pytest.mark.parametrize(
("name", "expr", "result"), data_test_bitwise_and_boolean_functions
)
def test_bit_and_bool_fns(df, name, expr, result):
df = df.aggregate([], [expr.alias(name)])
expected = {
name: result,
}
assert df.collect()[0].to_pydict() == expected
@pytest.mark.parametrize("fn", [f.bit_and, f.bit_or])
def test_bitwise_distinct_is_kept(fn):
# AND and OR ignore duplicates, so the result alone cannot show whether
# ``distinct`` reached the plan.
assert "DISTINCT" in fn(column("b"), distinct=True).canonical_name()
assert "DISTINCT" not in fn(column("b")).canonical_name()
@pytest.mark.parametrize(
("name", "expr", "result"),
[
("first_value", f.first_value(column("a")), [0, 4]),
(
"first_value_ordered",
f.first_value(column("a"), order_by=[column("a").sort(ascending=False)]),
[3, 6],
),
(
"first_value_with_null",
f.first_value(
column("b"),
order_by=[column("b").sort(ascending=True, nulls_first=True)],
null_treatment=NullTreatment.RESPECT_NULLS,
),
[None, None],
),
(
"first_value_no_list_order_by",
f.first_value(
column("b"),
order_by=column("b").sort(nulls_first=True),
null_treatment=NullTreatment.RESPECT_NULLS,
),
[None, None],
),
(
"first_value_ignore_null",
f.first_value(
column("b"),
order_by=[column("b").sort(ascending=True)],
null_treatment=NullTreatment.IGNORE_NULLS,
),
[7, 9],
),
(
"last_value_ordered",
f.last_value(column("a"), order_by=[column("a").sort(ascending=False)]),
[0, 4],
),
(
"last_value_no_list_ordered",
f.last_value(column("a"), order_by=column("a")),
[3, 6],
),
(
"last_value_with_null",
f.last_value(
column("b"),
order_by=[column("b").sort(ascending=True, nulls_first=False)],
null_treatment=NullTreatment.RESPECT_NULLS,
),
[None, None],
),
(
"last_value_ignore_null",
f.last_value(
column("b"),
order_by=[column("b").sort(ascending=True)],
null_treatment=NullTreatment.IGNORE_NULLS,
),
[8, 9],
),
(
"nth_value_ordered",
f.nth_value(column("a"), 2, order_by=[column("a").sort(ascending=False)]),
[2, 5],
),
(
"nth_value_no_list_ordered",
f.nth_value(column("a"), 2, order_by=column("a").sort(ascending=False)),
[2, 5],
),
(
"nth_value_with_null",
f.nth_value(
column("b"),
3,
order_by=[column("b").sort(ascending=True, nulls_first=False)],
null_treatment=NullTreatment.RESPECT_NULLS,
),
[8, None],
),
(
"nth_value_ignore_null",
f.nth_value(
column("b"),
2,
order_by=[column("b").sort(ascending=True)],
null_treatment=NullTreatment.IGNORE_NULLS,
),
[7, 9],
),
],
)
def test_first_last_value(df_partitioned, name, expr, result) -> None:
df = df_partitioned.aggregate([column("c")], [expr.alias(name)]).sort(column("c"))
expected = {
"c": ["A", "B"],
name: result,
}
assert df.collect()[0].to_pydict() == expected
@pytest.mark.parametrize(
("name", "expr", "result"),
[
("string_agg", f.string_agg(column("a"), ","), "one,two,three,two"),
("string_agg", f.string_agg(column("b"), ""), "03124"),
(
"string_agg",
f.string_agg(column("a"), ",", filter=column("b") != lit(3)),
"one,three,two",
),
(
"string_agg",
f.string_agg(column("a"), ",", order_by=[column("b")]),
"one,three,two,two",
),
(
"string_agg",
f.string_agg(column("a"), ",", order_by=column("b")),
"one,three,two,two",
),
(
"string_agg",
f.string_agg(column("a"), ",", distinct=True, order_by=column("a")),
"one,three,two",
),
],
)
def test_string_agg(name, expr, result) -> None:
ctx = SessionContext()
df = ctx.from_pydict(
{
"a": ["one", "two", None, "three", "two"],
"b": [0, 3, 1, 2, 4],
}
)
df = df.aggregate([], [expr.alias(name)])
expected = {
name: [result],
}
df.show()
assert df.collect()[0].to_pydict() == expected
_FILTER = column("b") > lit(1)
@pytest.mark.parametrize(
("call", "match"),
[
pytest.param(
lambda: f.percentile_cont(column("a"), 0.5, _FILTER),
r"percentile_cont\(\).*pass filter by keyword",
id="percentile_cont",
),
pytest.param(
lambda: f.quantile_cont(column("a"), 0.5, _FILTER),
r"quantile_cont\(\).*pass filter by keyword",
id="quantile_cont",
),
pytest.param(
lambda: f.mean(column("a"), _FILTER),
r"mean\(\).*pass filter by keyword",
id="mean",
),
pytest.param(
lambda: f.bit_and(column("a"), _FILTER),
r"bit_and\(\).*pass filter by keyword",
id="bit_and",
),
pytest.param(
lambda: f.bit_or(column("a"), _FILTER),
r"bit_or\(\).*pass filter by keyword",
id="bit_or",
),
pytest.param(
lambda: f.string_agg(column("a"), ",", _FILTER, column("b")),
r"string_agg\(\).*pass filter and order_by by keyword",
id="string_agg filter",
),
pytest.param(
lambda: f.string_agg(column("a"), ",", None, column("b")),
r"string_agg\(\).*pass filter and order_by by keyword",
id="string_agg None placeholder",
),
],
)
def test_positional_filter_names_the_function(call, match) -> None:
# ``distinct`` was inserted before ``filter``. A call written for the old
# signature must name the function and the fix, not fail inside PyO3. The
# ``None`` placeholder would otherwise run with order_by shifted into
# filter.
with pytest.raises(TypeError, match=match):
call()
def test_string_agg_accepts_numpy_bool_distinct() -> None:
np = pytest.importorskip("numpy")
df = SessionContext().from_pydict({"a": ["x", "y", "x"]})
expr = f.string_agg(column("a"), ",", distinct=np.True_, order_by="a")
assert df.aggregate([], [expr.alias("s")]).collect_column("s")[0].as_py() == "x,y"
@pytest.mark.parametrize(
("expr", "sql"),
[
pytest.param(
lambda s: f.percentile_cont(s, 0.25),
"percentile_cont(0.25)",
id="percentile_cont",
),
pytest.param(
lambda s: f.quantile_cont(s, 0.25),
"quantile_cont(0.25)",
id="quantile_cont",
),
pytest.param(
lambda s: f.approx_percentile_cont(s, 0.25),
"approx_percentile_cont(0.25)",
id="approx_percentile_cont",
),
pytest.param(
lambda s: f.approx_percentile_cont_with_weight(s, lit(1.0), 0.25),
"approx_percentile_cont_with_weight(1.0, 0.25)",
id="approx_percentile_cont_with_weight",
),
],
)
def test_percentile_keeps_sort_direction(expr, sql) -> None:
ctx = SessionContext()
df = ctx.from_pydict({"a": [1.0, 2.0, 3.0, 4.0, 5.0]}, name="t")
desc = column("a").sort(ascending=False)
result = df.aggregate([], [expr(desc).alias("p")]).collect_column("p")
expected = ctx.sql(
f"SELECT {sql} WITHIN GROUP (ORDER BY a DESC) AS p FROM t"
).collect_column("p")
assert result.to_pylist() == expected.to_pylist()
assert (
result.to_pylist()
!= df.aggregate([], [expr(column("a")).alias("p")])
.collect_column("p")
.to_pylist()
)
def test_percentile_cont_order_by_replaces_sort() -> None:
ctx = SessionContext()
df = ctx.from_pydict({"a": [1.0, 2.0, 3.0, 4.0, 5.0]})
expr = (
f.percentile_cont(column("a"), 0.25)
.order_by(column("a").sort(ascending=False))
.build()
)
assert df.aggregate([], [expr.alias("p")]).collect_column("p")[0].as_py() == 4.0