Add Count reduction operations (for groups).

docs/v0.1.0/.documenter-siteinfo.json

@@ -1 +1 @@
-{"documenter":{"julia_version":"1.10.2","generation_timestamp":"2024-04-02T14:36:25","documenter_version":"1.3.0"}}
+{"documenter":{"julia_version":"1.10.2","generation_timestamp":"2024-04-03T13:40:33","documenter_version":"1.3.0"}}

docs/v0.1.0/index.html

@@ -912,6 +912,12 @@ <h1 id="Index">
 </a>
 </li>
 <li>
+<a href="operations.html#Daf.Operations.Count">
+<code>Daf.Operations.Count
+</code>
+</a>
+</li>
+<li>
 <a href="operations.html#Daf.Operations.Fraction">
 <code>Daf.Operations.Fraction
 </code>

docs/v0.1.0/operations.html

@@ -541,6 +541,40 @@ <h2 id="Reduction-operations">
 <header>
 <a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring">
 </a>
+<a class="docstring-binding" id="Daf.Operations.Count" href="#Daf.Operations.Count">
+<code>Daf.Operations.Count
+</code>
+</a> — 
+<span class="docstring-category">Type
+</span>
+</header>
+<section>
+<div>
+<pre>
+<code class="language-julia hljs">Count(; dtype::Maybe{Type} = nothing)
+</code>
+</pre>
+<p>Reduction operation that counts elements. This is useful when using 
+<code>GroupBy
+</code> queries to count the number of elements in each group.
+</p>
+<p>
+<strong>Parameters
+</strong>
+</p>
+<p>
+<code>dtype
+</code> - By default, uses 
+<code>UInt32
+</code>.
+</p>
+</div>
+</section>
+</article>
+<article class="docstring">
+<header>
+<a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring">
+</a>
 <a class="docstring-binding" id="Daf.Operations.Sum" href="#Daf.Operations.Sum">
 <code>Daf.Operations.Sum
 </code>
@@ -1232,6 +1266,12 @@ <h2 id="Index">
 </a>
 </li>
 <li>
+<a href="operations.html#Daf.Operations.Count">
+<code>Daf.Operations.Count
+</code>
+</a>
+</li>
+<li>
 <a href="operations.html#Daf.Operations.Fraction">
 <code>Daf.Operations.Fraction
 </code>

src/operations.jl

@@ -26,6 +26,7 @@ using Daf.Generic
 export Abs
 export Clamp
 export Convert
+export Count
 export Fraction
 export Log
 export Max
@@ -744,7 +745,7 @@ end
 function compute_eltwise(operation::Significant, input::StorageMatrix{T})::StorageMatrix{T} where {T <: StorageNumber}
     output = copy(input)
     if output isa SparseMatrixCSC
-        @threads for column_index in 1:size(output)[2]
+        @threads for column_index in 1:size(output, 2)
             first = output.colptr[column_index]
             last = output.colptr[column_index + 1] - 1
             if first <= last
@@ -754,7 +755,7 @@ function compute_eltwise(operation::Significant, input::StorageMatrix{T})::Stora
         end
         dropzeros!(output)
     else
-        @threads for column_index in 1:size(output)[2]
+        @threads for column_index in 1:size(output, 2)
             column_vector = @view output[:, column_index]
             significant!(column_vector, operation.high, operation.low)
         end
@@ -793,6 +794,49 @@ function compute_eltwise(operation::Significant, input::T)::T where {T <: Storag
     return error("applying Significant eltwise operation to a scalar")
 end
 
+"""
+    Count(; dtype::Maybe{Type} = nothing)
+
+Reduction operation that counts elements. This is useful when using `GroupBy` queries to count the number of elements in
+each group.
+
+**Parameters**
+
+`dtype` - By default, uses `UInt32`.
+"""
+struct Count <: ReductionOperation
+    dtype::Maybe{Type}
+end
+@query_operation Count
+
+function Count(; dtype::Maybe{Type} = nothing)::Count
+    @assert dtype == nothing || dtype <: Real
+    return Count(dtype)
+end
+
+function Count(operation_name::Token, parameters_values::Dict{String, Token})::Count
+    dtype = parse_parameter_value(operation_name, "reduction", parameters_values, "dtype", nothing) do parameter_value
+        return parse_number_dtype_value(operation_name, "dtype", parameter_value)
+    end
+    return Count(dtype)
+end
+
+function compute_reduction(operation::Count, input::StorageMatrix{T})::StorageVector where {T <: StorageNumber}
+    dtype = reduction_result_type(operation, T)
+    result = Vector{dtype}(undef, size(input, 2))
+    result .= size(input, 1)
+    return result
+end
+
+function compute_reduction(operation::Count, input::StorageVector{T})::StorageNumber where {T <: StorageNumber}
+    dtype = reduction_result_type(operation, T)
+    return dtype(length(input))
+end
+
+function reduction_result_type(operation::Count, eltype::Type)::Type
+    return operation.dtype == nothing ? UInt32 : operation.dtype
+end
+
 """
     Sum(; dtype::Maybe{Type} = nothing)
 
@@ -820,7 +864,7 @@ end
 
 function compute_reduction(operation::Sum, input::StorageMatrix{T})::StorageVector where {T <: StorageNumber}
     dtype = reduction_result_type(operation, T)
-    result = Vector{dtype}(undef, size(input)[2])
+    result = Vector{dtype}(undef, size(input, 2))
     sum!(transpose(result), input)
     return result
 end
@@ -963,7 +1007,7 @@ end
 
 function compute_reduction(operation::Quantile, input::StorageMatrix{T})::StorageVector where {T <: StorageNumber}
     dtype = reduction_result_type(operation, eltype(input))
-    output = Vector{dtype}(undef, size(input)[2])
+    output = Vector{dtype}(undef, size(input, 2))
     @threads for column_index in 1:length(output)
         column_vector = @view input[:, column_index]
         output[column_index] = quantile(column_vector, operation.p)  # NOJET

src/operations.md

@@ -19,6 +19,7 @@ Daf.Operations.Significant
 ## Reduction operations
 
 ```@docs
+Daf.Operations.Count
 Daf.Operations.Sum
 Daf.Operations.Max
 Daf.Operations.Min

test/operations.jl

@@ -285,6 +285,19 @@ nested_test("operations") do
     end
 
     nested_test("reduction") do
+        nested_test("count") do
+            nested_test("vector") do
+                set_vector!(daf, "cell", "value", [1, 3])
+                @test with_type(daf[Axis("cell") |> Lookup("value") |> Count()]) == (2, UInt32)
+                @test with_type(daf["/ cell : value %> Count dtype Int8"]) == (2, Int8)
+            end
+
+            nested_test("matrix") do
+                set_matrix!(daf, "cell", "gene", "value", [1.0 2.0 2.0; -3.0 1.0 6.0])
+                @test with_type(daf["/ cell / gene : value %> Count"]) == ([2, 2, 2], UInt32)
+            end
+        end
+
         nested_test("sum") do
             nested_test("vector") do
                 set_vector!(daf, "cell", "value", [1, 3])

test/queries.jl

@@ -1015,6 +1015,10 @@ nested_test("queries") do
                 @test get_result(daf, q"/ cell : age @ batch %> Min") == ("batch", ["U" => 1, "V" => 2])
             end
 
+            nested_test("()") do
+                @test get_result(daf, q"/ cell : age @ batch %> Count") == ("batch", ["U" => 2, "V" => 1])
+            end
+
             nested_test("!reduction") do
                 return test_invalid(
                     daf,
@@ -1051,6 +1055,20 @@ nested_test("queries") do
                 )
             end
 
+            nested_test("()") do
+                @test get_result(daf, q"/ cell / gene : UMIs @ batch %> Count") == (
+                    ("batch", "gene"),
+                    [
+                        ("U", "X") => 2,
+                        ("U", "Y") => 2,
+                        ("U", "Z") => 2,
+                        ("V", "X") => 1,
+                        ("V", "Y") => 1,
+                        ("V", "Z") => 1,
+                    ],
+                )
+            end
+
             nested_test("!reduction") do
                 return test_invalid(
                     daf,