Add the count_by, aggregate_by and chunk operators

Ports Enumerable.CountBy and AggregateBy (.NET 9) and Enumerable.Chunk
(.NET 6).

count_by and aggregate_by name their key columns through the existing
_key_namedtuple, so they agree with summarize: a scalar key becomes a column
called `key`, a NamedTuple key contributes one column per field. A test
asserts count_by and groupby-plus-summarize produce identical output.

summarize remains the more general and more idiomatic way to aggregate;
aggregate_by is here for LINQ parity, and its accumulator takes
(accumulated, element) to match .NET's argument order.

chunk pulls at most n elements per call, so a chunked source is only walked
as far as the batches actually consumed.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: davidanthoff

src/QueryOperators.jl

@@ -36,6 +36,9 @@ include("enumerable/enumerable_shuffle.jl")
 include("enumerable/enumerable_index.jl")
 include("enumerable/enumerable_pivot.jl")
 include("enumerable/enumerable_summarize.jl")
+include("enumerable/enumerable_countby.jl")
+include("enumerable/enumerable_aggregateby.jl")
+include("enumerable/enumerable_chunk.jl")
 include("enumerable/show.jl")
 
 include("source_iterable.jl")

src/enumerable/enumerable_aggregateby.jl

@@ -0,0 +1,49 @@
+_aggregate_by_row(k, v) = Base.merge(_key_namedtuple(k), (value = v,))
+
+struct EnumerableAggregateBy{T,TKey,TACC,S,Q<:Function,A<:Function} <: Enumerable
+    source::S
+    f::Q
+    seed::TACC
+    accumulator::A
+end
+
+Base.eltype(::Type{EnumerableAggregateBy{T,TKey,TACC,S,Q,A}}) where {T,TKey,TACC,S,Q,A} = T
+
+# Enumerable.AggregateBy (.NET 9): fold the elements of each key into a single
+# value, without materialising the intermediate groupings. `accumulator` is
+# called as `accumulator(accumulated, element)`, matching .NET's argument order.
+#
+# `summarize` is the more general and more idiomatic way to aggregate here;
+# `aggregate_by` exists for LINQ parity.
+function aggregate_by(source::Enumerable, f::Function, f_expr::Expr, seed, accumulator::Function)
+    TS = eltype(source)
+    TKey = Base._return_type(f, Tuple{TS,})
+    TACC = typeof(seed)
+    T = Base._return_type(_aggregate_by_row, Tuple{TKey,TACC})
+
+    return EnumerableAggregateBy{T,TKey,TACC,typeof(source),typeof(f),typeof(accumulator)}(source, f, seed, accumulator)
+end
+
+function Base.iterate(iter::EnumerableAggregateBy{T,TKey,TACC,S,Q,A}) where {T,TKey,TACC,S,Q,A}
+    accumulated = OrderedDict{TKey,TACC}()
+    for i in iter.source
+        k = iter.f(i)
+        accumulated[k] = iter.accumulator(get(accumulated, k, iter.seed), i)
+    end
+
+    rows = T[_aggregate_by_row(k, v) for (k, v) in accumulated]
+
+    if length(rows)==0
+        return nothing
+    end
+
+    return rows[1], (rows, 2)
+end
+
+function Base.iterate(iter::EnumerableAggregateBy{T,TKey,TACC,S,Q,A}, state) where {T,TKey,TACC,S,Q,A}
+    if state[2]>length(state[1])
+        return nothing
+    else
+        return state[1][state[2]], (state[1], state[2]+1)
+    end
+end

src/enumerable/enumerable_chunk.jl

@@ -0,0 +1,42 @@
+struct EnumerableChunk{T,TI,S} <: Enumerable
+    source::S
+    n::Int
+end
+
+# Enumerable.Chunk (.NET 6): split the source into batches of at most `n`
+# elements. The final batch is shorter when the source does not divide evenly.
+function chunk(source::Enumerable, n::Integer)
+    n < 1 && error("The chunk size must be at least 1, got $n.")
+
+    TI = eltype(source)
+    T = Vector{TI}
+
+    return EnumerableChunk{T,TI,typeof(source)}(source, Int(n))
+end
+
+Base.IteratorSize(::Type{EnumerableChunk{T,TI,S}}) where {T,TI,S} = haslength(S)
+
+Base.eltype(::Type{EnumerableChunk{T,TI,S}}) where {T,TI,S} = T
+
+Base.length(iter::EnumerableChunk) = cld(length(iter.source), iter.n)
+
+Base.iterate(iter::EnumerableChunk) = _chunk_next(iter, _NotStarted())
+
+Base.iterate(iter::EnumerableChunk, state) = _chunk_next(iter, state)
+
+# Pulls at most `n` elements per call, so a chunked source is only walked as
+# far as the batches actually consumed.
+function _chunk_next(iter::EnumerableChunk{T,TI,S}, source_state) where {T,TI,S}
+    buffer = TI[]
+
+    while length(buffer) < iter.n
+        ret = _iterate_from(iter.source, source_state)
+        ret === nothing && break
+        push!(buffer, ret[1])
+        source_state = ret[2]
+    end
+
+    length(buffer)==0 && return nothing
+
+    return buffer, source_state
+end

src/enumerable/enumerable_countby.jl

@@ -0,0 +1,44 @@
+# Key columns are named exactly as `summarize` names them: a scalar key becomes
+# a column called `key`, a NamedTuple key contributes one column per field.
+_count_by_row(k, n::Int) = Base.merge(_key_namedtuple(k), (count = n,))
+
+struct EnumerableCountBy{T,TKey,S,Q<:Function} <: Enumerable
+    source::S
+    f::Q
+end
+
+Base.eltype(::Type{EnumerableCountBy{T,TKey,S,Q}}) where {T,TKey,S,Q} = T
+
+# Enumerable.CountBy (.NET 9): the frequency of each key, without materialising
+# the intermediate groupings that `groupby` would build.
+function count_by(source::Enumerable, f::Function, f_expr::Expr)
+    TS = eltype(source)
+    TKey = Base._return_type(f, Tuple{TS,})
+    T = Base._return_type(_count_by_row, Tuple{TKey,Int})
+
+    return EnumerableCountBy{T,TKey,typeof(source),typeof(f)}(source, f)
+end
+
+function Base.iterate(iter::EnumerableCountBy{T,TKey,S,Q}) where {T,TKey,S,Q}
+    counts = OrderedDict{TKey,Int}()
+    for i in iter.source
+        k = iter.f(i)
+        counts[k] = get(counts, k, 0) + 1
+    end
+
+    rows = T[_count_by_row(k, n) for (k, n) in counts]
+
+    if length(rows)==0
+        return nothing
+    end
+
+    return rows[1], (rows, 2)
+end
+
+function Base.iterate(iter::EnumerableCountBy{T,TKey,S,Q}, state) where {T,TKey,S,Q}
+    if state[2]>length(state[1])
+        return nothing
+    else
+        return state[1][state[2]], (state[1], state[2]+1)
+    end
+end

src/operators.jl

@@ -170,3 +170,13 @@ function reverse end
 function shuffle end
 
 function index end
+
+# Keyed aggregation and batching. `count_by` and `aggregate_by` are
+# Enumerable.CountBy and AggregateBy (.NET 9); `chunk` is Enumerable.Chunk
+# (.NET 6).
+
+function count_by end
+
+function aggregate_by end
+
+function chunk end

test/runtests.jl

@@ -4,6 +4,7 @@ include("test_core.jl")
 include("test_enumerable_outerjoins.jl")
 include("test_enumerable_setops.jl")
 include("test_enumerable_ordering.jl")
+include("test_enumerable_keyed_aggregation.jl")
 include("test_enumerable_unique.jl")
 include("test_enumerable_summarize.jl")
 include("test_namedtupleutilities.jl")

test/test_enumerable_keyed_aggregation.jl

@@ -0,0 +1,165 @@
+@testitem "count_by with a scalar key" begin
+    source = QueryOperators.query([(k="a", v=1), (k="b", v=2), (k="a", v=3)])
+
+    res = QueryOperators.count_by(source, i -> i.k, :(i -> i.k))
+
+    # A scalar key becomes a column called `key`, as summarize names it.
+    @test collect(res) == [(key="a", count=2), (key="b", count=1)]
+    @test eltype(res) == NamedTuple{(:key, :count),Tuple{String,Int}}
+end
+
+@testitem "count_by with a NamedTuple key splats the key columns" begin
+    source = QueryOperators.query([(a=1, b=1), (a=1, b=1), (a=1, b=2)])
+
+    res = QueryOperators.count_by(source, i -> (a=i.a, b=i.b), :(i -> (a=i.a, b=i.b)))
+
+    @test collect(res) == [(a=1, b=1, count=2), (a=1, b=2, count=1)]
+end
+
+@testitem "count_by preserves first-seen key order" begin
+    source = QueryOperators.query([3, 1, 3, 2, 1])
+
+    res = QueryOperators.count_by(source, i -> i, :(i -> i))
+
+    @test collect(res) == [(key=3, count=2), (key=1, count=2), (key=2, count=1)]
+end
+
+@testitem "count_by on an empty source" begin
+    source = QueryOperators.query(Int[])
+
+    @test collect(QueryOperators.count_by(source, i -> i, :(i -> i))) == NamedTuple{(:key, :count),Tuple{Int,Int}}[]
+end
+
+@testitem "count_by agrees with groupby plus summarize" begin
+    using Statistics
+
+    data = [(k=1, v=10), (k=2, v=20), (k=1, v=30), (k=3, v=40), (k=1, v=50)]
+
+    by_count = collect(QueryOperators.count_by(QueryOperators.query(data), i -> i.k, :(i -> i.k)))
+
+    grouped = QueryOperators.@groupby_simple(QueryOperators.query(data), i -> i.k)
+    by_summarize = collect(QueryOperators.summarize(
+        grouped,
+        g -> Base.merge(QueryOperators._key_namedtuple(QueryOperators.key(g)), (count = length(g),)),
+        :(g -> (count = length(g),))))
+
+    @test by_count == by_summarize
+end
+
+@testitem "aggregate_by" begin
+    source = QueryOperators.query([(id="0", score=42), (id="1", score=5), (id="2", score=4), (id="1", score=10), (id="0", score=25)])
+
+    res = QueryOperators.aggregate_by(
+        source,
+        i -> i.id, :(i -> i.id),
+        0,
+        (total, cur) -> total + cur.score)
+
+    # The example from the .NET 9 release notes.
+    @test collect(res) == [(key="0", value=67), (key="1", value=15), (key="2", value=4)]
+    @test eltype(res) == NamedTuple{(:key, :value),Tuple{String,Int}}
+end
+
+@testitem "aggregate_by with a NamedTuple key" begin
+    source = QueryOperators.query([(a=1, b=1, v=2), (a=1, b=1, v=3), (a=2, b=1, v=5)])
+
+    res = QueryOperators.aggregate_by(
+        source,
+        i -> (a=i.a, b=i.b), :(i -> (a=i.a, b=i.b)),
+        1,
+        (acc, cur) -> acc * cur.v)
+
+    @test collect(res) == [(a=1, b=1, value=6), (a=2, b=1, value=5)]
+end
+
+@testitem "aggregate_by on an empty source" begin
+    source = QueryOperators.query(Int[])
+
+    res = QueryOperators.aggregate_by(source, i -> i, :(i -> i), 0, (acc, cur) -> acc + cur)
+
+    @test collect(res) == NamedTuple{(:key, :value),Tuple{Int,Int}}[]
+end
+
+@testitem "aggregate_by never shares the seed between keys" begin
+    source = QueryOperators.query([1, 1, 2, 2, 2])
+
+    res = QueryOperators.aggregate_by(source, i -> i, :(i -> i), Int[], (acc, cur) -> vcat(acc, cur))
+
+    @test collect(res) == [(key=1, value=[1, 1]), (key=2, value=[2, 2, 2])]
+end
+
+@testitem "chunk" begin
+    source = QueryOperators.query([1, 2, 3, 4, 5])
+
+    res = QueryOperators.chunk(source, 2)
+
+    # The final chunk is short when the source does not divide evenly.
+    @test collect(res) == [[1, 2], [3, 4], [5]]
+    @test eltype(res) == Vector{Int}
+    @test Base.IteratorSize(typeof(res)) == Base.HasLength()
+    @test length(res) == 3
+end
+
+@testitem "chunk when the size divides evenly" begin
+    source = QueryOperators.query([1, 2, 3, 4])
+
+    @test collect(QueryOperators.chunk(source, 2)) == [[1, 2], [3, 4]]
+    @test length(QueryOperators.chunk(source, 2)) == 2
+end
+
+@testitem "chunk with a size larger than the source" begin
+    source = QueryOperators.query([1, 2])
+
+    @test collect(QueryOperators.chunk(source, 10)) == [[1, 2]]
+end
+
+@testitem "chunk on an empty source" begin
+    source = QueryOperators.query(Int[])
+
+    @test collect(QueryOperators.chunk(source, 3)) == Vector{Int}[]
+end
+
+@testitem "chunk rejects a size below 1" begin
+    source = QueryOperators.query([1, 2, 3])
+
+    @test_throws ErrorException QueryOperators.chunk(source, 0)
+    @test_throws ErrorException QueryOperators.chunk(source, -1)
+end
+
+@testitem "chunk only walks the source as far as the batches consumed" begin
+    import IteratorInterfaceExtensions
+
+    mutable struct CountedChunkSource
+        data::Vector{Int}
+        pulled::Int
+    end
+    Base.eltype(::Type{CountedChunkSource}) = Int
+    Base.IteratorSize(::Type{CountedChunkSource}) = Base.HasLength()
+    Base.length(c::CountedChunkSource) = length(c.data)
+    function Base.iterate(c::CountedChunkSource, i=1)
+        i > length(c.data) && return nothing
+        c.pulled += 1
+        return c.data[i], i + 1
+    end
+    IteratorInterfaceExtensions.isiterable(::CountedChunkSource) = true
+    IteratorInterfaceExtensions.getiterator(c::CountedChunkSource) = c
+
+    src = CountedChunkSource(collect(1:100), 0)
+    res = QueryOperators.chunk(QueryOperators.query(src), 3)
+
+    it = iterate(res)
+
+    @test it[1] == [1, 2, 3]
+    @test src.pulled == 3
+end
+
+@testitem "chunk works downstream of groupby" begin
+    source = QueryOperators.query([(k=1, v=1), (k=2, v=2), (k=3, v=3)])
+    grouped = QueryOperators.@groupby_simple(source, i -> i.k)
+
+    chunks = collect(QueryOperators.chunk(grouped, 2))
+
+    @test length(chunks) == 2
+    @test [QueryOperators.key(g) for g in chunks[1]] == [1, 2]
+    @test [QueryOperators.key(g) for g in chunks[2]] == [3]
+end