(beam) branch master updated: Return zero elements immediately if the requested number of quantiles is 1. (#33524)

[lostluck]

This is an automated email from the ASF dual-hosted git repository.

lostluck pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 6618968a2f7 Return zero elements immediately if the requested number 
of quantiles is 1. (#33524)
6618968a2f7 is described below

commit 6618968a2f7a0548915161259dfc2dd9bdb002b5
Author: Dmitry Labutin <labu...@google.com>
AuthorDate: Tue Jan 7 13:00:49 2025 -0800

    Return zero elements immediately if the requested number of quantiles is 1. 
(#33524)
---
 sdks/go/pkg/beam/transforms/stats/quantiles.go      |  4 ++++
 sdks/go/pkg/beam/transforms/stats/quantiles_test.go | 17 +++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/sdks/go/pkg/beam/transforms/stats/quantiles.go 
b/sdks/go/pkg/beam/transforms/stats/quantiles.go
index 6d2baa8b5e9..93ca4e3ee52 100644
--- a/sdks/go/pkg/beam/transforms/stats/quantiles.go
+++ b/sdks/go/pkg/beam/transforms/stats/quantiles.go
@@ -701,6 +701,10 @@ func reduce(s beam.Scope, weightedElements 
beam.PCollection, state approximateQu
 // For example, if numQuantiles = 2, the returned list would contain a single 
element such that approximately half of the input would be less than that 
element and half would be greater or equal.
 func ApproximateWeightedQuantiles(s beam.Scope, pc beam.PCollection, less any, 
opts Opts) beam.PCollection {
        _, t := beam.ValidateKVType(pc)
+       // Return zero elements immediately if the requested number of 
quantiles is 1.
+       if opts.NumQuantiles == 1 {
+               return beam.Create(s, 
reflect.New(reflect.SliceOf(t.Type())).Elem().Interface())
+       }
        state := approximateQuantilesCombineFnState{
                K:            opts.K,
                NumQuantiles: opts.NumQuantiles,
diff --git a/sdks/go/pkg/beam/transforms/stats/quantiles_test.go 
b/sdks/go/pkg/beam/transforms/stats/quantiles_test.go
index 1e389eed128..bad5fce3062 100644
--- a/sdks/go/pkg/beam/transforms/stats/quantiles_test.go
+++ b/sdks/go/pkg/beam/transforms/stats/quantiles_test.go
@@ -246,3 +246,20 @@ func TestWeightedElementEncoding(t *testing.T) {
                t.Errorf("Invalid coder. Wanted %v got %v", w, decoded)
        }
 }
+
+func TestZeroQuantiles(t *testing.T) {
+       const numElements int = 30000
+       inputSlice := make([]int, 0, numElements)
+       for i := 0; i < numElements; i++ {
+               inputSlice = append(inputSlice, i)
+       }
+       p, s, input, expected := ptest.CreateList2(inputSlice, [][]int{{}})
+       quantiles := ApproximateQuantiles(s, input, less, Opts{
+               K:            200,
+               NumQuantiles: 1,
+       })
+       passert.Equals(s, quantiles, expected)
+       if err := ptest.Run(p); err != nil {
+               t.Errorf("ApproximateQuantiles failed: %v", err)
+       }
+}