feat: accept distinct kwarg on sum and avg (#1556)

* feat: accept distinct kwarg on sum and avg

Upstream exposes `sum_distinct` / `avg_distinct` / `count_distinct` as
sibling functions that call the same underlying UDAF with
`distinct: bool = true`. The Rust binding side already routes
`distinct=Some(true)` through the aggregate builder for `sum`, `avg`,
and `count` — but only `count` exposed the kwarg on the Python wrapper.

Add `distinct: bool = False` to `sum()` and `avg()` mirroring the
existing `count()` signature, and update SKILL.md so the check-upstream
audit does not re-flag the three upstream `*_distinct` shortcuts as
gaps. The plan emitted by `sum(col, distinct=True)` matches what
upstream's `sum_distinct(col)` builds.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* test: fold sum/avg distinct tests into parameterized aggregation test

Move the standalone test_sum_distinct_kwarg and test_avg_distinct_kwarg
from test_functions.py into the existing test_aggregation::test_aggregation
parameterization, matching how distinct is already covered for median,
array_agg, count, and bit_xor.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* docs: clarify distinct kwarg on sum and avg

Drop the unhelpful "upstream avg_distinct/sum_distinct shortcut"
reference in favor of describing the actual behavior.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* docs: note sum/avg distinct argument-order breaking change

distinct is inserted before filter on sum and avg for consistency with
the other aggregate functions, breaking positional filter callers. Add a
DataFusion 54.0.0 upgrade-guide entry covering the migration.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Update docs/source/user-guide/upgrade-guides.rst

Co-authored-by: Nick <24689722+ntjohnson1@users.noreply.github.com>

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Nick <24689722+ntjohnson1@users.noreply.github.com>

Author: 3 people

.ai/skills/check-upstream/SKILL.md

@@ -88,11 +88,18 @@ The user may specify an area via `$ARGUMENTS`. If no area is specified or "all"
 - Python API: `python/datafusion/functions.py` (aggregate functions are mixed in with scalar functions)
 - Rust bindings: `crates/core/src/functions.rs`
 
+**Evaluated and not requiring separate Python exposure:**
+- `count_distinct` — covered by `count(expr, distinct=True)`. Both forms call
+  `count_udaf` with `distinct: bool = true` and produce the same logical plan.
+- `sum_distinct` — covered by `sum(expr, distinct=True)`.
+- `avg_distinct` — covered by `avg(expr, distinct=True)`.
+
 **How to check:**
 1. Fetch the upstream aggregate function documentation page
 2. Compare against aggregate functions in `python/datafusion/functions.py` (check `__all__` list and function definitions)
 3. A function is covered if it exists in the Python API, even if it aliases another function's Rust binding
-4. Report only functions missing from the Python API
+4. Check against the "evaluated and not requiring exposure" list before flagging as a gap
+5. Report only functions missing from the Python API
 
 ### 3. Window Functions
 

docs/source/user-guide/upgrade-guides.rst

@@ -45,6 +45,26 @@ After:
     config = SessionConfig().set("datafusion.execution.batch_size", "4096")
     ctx = SessionContext(config)
 
+The aggregate functions :py:func:`~datafusion.functions.sum` and
+:py:func:`~datafusion.functions.avg` now accept a ``distinct`` argument, matching
+the other aggregate functions. ``distinct`` is inserted *before* ``filter`` in the
+argument list, so any code that passed ``filter`` positionally must be updated to
+pass it as a keyword argument. The types are distinct so a type checker should flag this.
+
+Before:
+
+.. code-block:: python
+
+    f.sum(column("a"), my_filter)
+    f.avg(column("a"), my_filter)
+
+Now:
+
+.. code-block:: python
+
+    f.sum(column("a"), filter=my_filter)
+    f.avg(column("a"), filter=my_filter)
+
 DataFusion 53.0.0
 -----------------
 

python/datafusion/functions.py

@@ -4563,17 +4563,19 @@ def grouping(
 
 def avg(
     expression: Expr,
+    distinct: bool = False,
     filter: Expr | None = None,
 ) -> Expr:
     """Returns the average value.
 
     This aggregate function expects a numeric expression and will return a float.
 
     If using the builder functions described in ref:`_aggregation` this function ignores
-    the options ``order_by``, ``null_treatment``, and ``distinct``.
+    the options ``order_by`` and ``null_treatment``.
 
     Args:
         expression: Values to combine into an array
+        distinct: If True, duplicate values are removed before averaging.
         filter: If provided, only compute against rows for which the filter is True
 
     Examples:
@@ -4593,9 +4595,17 @@ def avg(
         ...     ).alias("v")])
         >>> result.collect_column("v")[0].as_py()
         2.5
+
+        >>> df = ctx.from_pydict({"a": [1.0, 1.0, 2.0, 3.0]})
+        >>> result = df.aggregate(
+        ...     [], [dfn.functions.avg(
+        ...         dfn.col("a"), distinct=True,
+        ...     ).alias("v")])
+        >>> result.collect_column("v")[0].as_py()
+        2.0
     """
     filter_raw = filter.expr if filter is not None else None
-    return Expr(f.avg(expression.expr, filter=filter_raw))
+    return Expr(f.avg(expression.expr, distinct=distinct, filter=filter_raw))
 
 
 def corr(value_y: Expr, value_x: Expr, filter: Expr | None = None) -> Expr:
@@ -4880,17 +4890,19 @@ def min(expression: Expr, filter: Expr | None = None) -> Expr:
 
 def sum(
     expression: Expr,
+    distinct: bool = False,
     filter: Expr | None = None,
 ) -> Expr:
     """Computes the sum of a set of numbers.
 
     This aggregate function expects a numeric expression.
 
     If using the builder functions described in ref:`_aggregation` this function ignores
-    the options ``order_by``, ``null_treatment``, and ``distinct``.
+    the options ``order_by`` and ``null_treatment``.
 
     Args:
         expression: Values to combine into an array
+        distinct: If True, duplicate values are removed before summing.
         filter: If provided, only compute against rows for which the filter is True
 
     Examples:
@@ -4910,9 +4922,17 @@ def sum(
         ...     ).alias("v")])
         >>> result.collect_column("v")[0].as_py()
         5
+
+        >>> df = ctx.from_pydict({"a": [1, 1, 2, 3]})
+        >>> result = df.aggregate(
+        ...     [], [dfn.functions.sum(
+        ...         dfn.col("a"), distinct=True,
+        ...     ).alias("v")])
+        >>> result.collect_column("v")[0].as_py()
+        6
     """
     filter_raw = filter.expr if filter is not None else None
-    return Expr(f.sum(expression.expr, filter=filter_raw))
+    return Expr(f.sum(expression.expr, distinct=distinct, filter=filter_raw))
 
 
 def stddev(expression: Expr, filter: Expr | None = None) -> Expr:

python/tests/test_aggregation.py

@@ -192,7 +192,9 @@ def test_aggregation_stats(df, agg_expr, calc_expected):
             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),