From 871f17e4d4be0f6995fb5688e48200bd4c858d15 Mon Sep 17 00:00:00 2001 From: Christopher Krah <48338417+chrikrah@users.noreply.github.com> Date: Thu, 5 Nov 2026 08:15:37 +0100 Subject: [PATCH] GH-55372: [C++][Python] Reject q outside [0, 1] in tdigest and hash_tdigest Signed-off-by: Christopher Krah <48338417+chrikrah@users.noreply.github.com> --- cpp/src/arrow/acero/hash_aggregate_test.cc | 18 ++++++++++ .../compute/kernels/aggregate_internal.h | 12 +++++++ .../compute/kernels/aggregate_quantile.cc | 8 ++--- .../compute/kernels/aggregate_tdigest.cc | 5 +-- .../arrow/compute/kernels/aggregate_test.cc | 36 +++++++++++++++++++ .../compute/kernels/hash_aggregate_numeric.cc | 1 + python/pyarrow/tests/test_compute.py | 13 +++++-- 7 files changed, 83 insertions(+), 10 deletions(-) diff --git a/cpp/src/arrow/acero/hash_aggregate_test.cc b/cpp/src/arrow/acero/hash_aggregate_test.cc index 2f498f3a7bb2..f30c48cf0032 100644 --- a/cpp/src/arrow/acero/hash_aggregate_test.cc +++ b/cpp/src/arrow/acero/hash_aggregate_test.cc @@ -1316,6 +1316,24 @@ TEST_P(GroupBy, TDigest) { /*verbose=*/true); } +TEST_P(GroupBy, TDigestInvalidQuantileLevel) { + auto argument = ArrayFromJSON(float64(), "[1.0, 2.0, 3.0]"); + auto key = ArrayFromJSON(int64(), "[1, 1, 2]"); + auto empty_argument = ArrayFromJSON(float64(), "[]"); + auto empty_key = ArrayFromJSON(int64(), "[]"); + for (double q : {-0.1, 1.1, std::numeric_limits::quiet_NaN()}) { + auto options = std::make_shared(q); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, HasSubstr("Quantile must be between 0 and 1"), + GroupByTest({argument}, {key}, {{"hash_tdigest", options}}, + /*use_threads=*/false)); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, HasSubstr("Quantile must be between 0 and 1"), + GroupByTest({empty_argument}, {empty_key}, {{"hash_tdigest", options}}, + /*use_threads=*/false)); + } +} + TEST_P(GroupBy, TDigestDecimal) { auto batch = RecordBatchFromJSON( schema({field("argument0", decimal128(3, 2)), field("argument1", decimal256(3, 2)), diff --git a/cpp/src/arrow/compute/kernels/aggregate_internal.h b/cpp/src/arrow/compute/kernels/aggregate_internal.h index 126f2891e212..67a427859f3a 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_internal.h +++ b/cpp/src/arrow/compute/kernels/aggregate_internal.h @@ -19,8 +19,10 @@ #include #include +#include #include "arrow/compute/kernels/util_internal.h" +#include "arrow/status.h" #include "arrow/type.h" #include "arrow/type_traits.h" #include "arrow/util/bit_run_reader.h" @@ -29,6 +31,16 @@ namespace arrow::compute::internal { +// Quantile levels must lie in [0, 1]. The negated form also rejects NaN. +inline Status ValidateQuantileLevels(const std::vector& q) { + for (double level : q) { + if (!(level >= 0 && level <= 1)) { + return Status::Invalid("Quantile must be between 0 and 1"); + } + } + return Status::OK(); +} + // Find the largest compatible primitive type for a primitive type. template struct FindAccumulatorType {}; diff --git a/cpp/src/arrow/compute/kernels/aggregate_quantile.cc b/cpp/src/arrow/compute/kernels/aggregate_quantile.cc index 79b9b415df40..2380ce4eafb7 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_quantile.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_quantile.cc @@ -20,6 +20,7 @@ #include #include "arrow/compute/api_aggregate.h" +#include "arrow/compute/kernels/aggregate_internal.h" #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/kernels/util_internal.h" #include "arrow/compute/registry_internal.h" @@ -488,12 +489,7 @@ Status CheckQuantileOptions(KernelContext* ctx) { if (options.q.empty()) { return Status::Invalid("Requires quantile argument"); } - for (double q : options.q) { - if (q < 0 || q > 1) { - return Status::Invalid("Quantile must be between 0 and 1"); - } - } - return Status::OK(); + return ValidateQuantileLevels(options.q); } template diff --git a/cpp/src/arrow/compute/kernels/aggregate_tdigest.cc b/cpp/src/arrow/compute/kernels/aggregate_tdigest.cc index 7ebc85941135..8ac003f4f7fe 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_tdigest.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_tdigest.cc @@ -167,8 +167,9 @@ struct TDigestInitState { Result> TDigestInit(KernelContext* ctx, const KernelInitArgs& args) { - TDigestInitState visitor(ctx, *args.inputs[0].type, - static_cast(*args.options)); + const auto& options = static_cast(*args.options); + RETURN_NOT_OK(ValidateQuantileLevels(options.q)); + TDigestInitState visitor(ctx, *args.inputs[0].type, options); return visitor.Create(); } diff --git a/cpp/src/arrow/compute/kernels/aggregate_test.cc b/cpp/src/arrow/compute/kernels/aggregate_test.cc index 4ef2b2f88a51..4ecdc11a7a29 100644 --- a/cpp/src/arrow/compute/kernels/aggregate_test.cc +++ b/cpp/src/arrow/compute/kernels/aggregate_test.cc @@ -4399,6 +4399,20 @@ TEST(TestQuantileKernel, Scalar) { } } +TEST(TestQuantileKernel, InvalidQuantileLevel) { + const double nan = std::numeric_limits::quiet_NaN(); + auto values = ArrayFromJSON(float64(), "[1.0, 2.0, 3.0]"); + auto empty = ArrayFromJSON(float64(), "[]"); + for (const std::vector& q : + std::vector>{{-0.1}, {1.1}, {nan}, {0.5, nan}}) { + for (const auto& input : {values, empty}) { + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("Quantile must be between 0 and 1"), + Quantile(input, QuantileOptions(q))); + } + } +} + TEST(TestQuantileKernel, Options) { auto ty = float64(); QuantileOptions keep_nulls(/*q=*/0.5, QuantileOptions::LINEAR, @@ -4516,6 +4530,28 @@ TEST(TestTDigestKernel, Options) { ResultWith(ArrayFromJSON(ty, "[null]"))); } +TEST(TestTDigestKernel, InvalidQuantileLevel) { + const double nan = std::numeric_limits::quiet_NaN(); + auto values = ArrayFromJSON(float64(), "[1.0, 2.0, 3.0]"); + auto empty = ArrayFromJSON(float64(), "[]"); + for (const std::vector& q : + std::vector>{{-0.1}, + {1.1}, + {0.5, 1.5}, + {-std::numeric_limits::infinity()}, + {nan}, + {0.5, nan}}) { + for (const auto& input : {values, empty}) { + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("Quantile must be between 0 and 1"), + TDigest(input, TDigestOptions(q))); + } + } + // The bounds themselves are valid. + EXPECT_THAT(TDigest(values, TDigestOptions(std::vector{0, 1})), + ResultWith(ArrayFromJSON(float64(), "[1.0, 3.0]"))); +} + TEST(TestTDigestKernel, ApproximateMedian) { // This is a wrapper for TDigest for (const auto& ty : {float64(), int64(), uint16()}) { diff --git a/cpp/src/arrow/compute/kernels/hash_aggregate_numeric.cc b/cpp/src/arrow/compute/kernels/hash_aggregate_numeric.cc index 30504ba87b6b..8fd64870d46c 100644 --- a/cpp/src/arrow/compute/kernels/hash_aggregate_numeric.cc +++ b/cpp/src/arrow/compute/kernels/hash_aggregate_numeric.cc @@ -924,6 +924,7 @@ struct GroupedTDigestImpl : public GroupedAggregator { Status Init(ExecContext* ctx, const KernelInitArgs& args) override { options_ = *checked_cast(args.options); + RETURN_NOT_OK(ValidateQuantileLevels(options_.q)); if (is_decimal_type::value) { decimal_scale_ = checked_cast(*args.inputs[0].type).scale(); } else { diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index 691d2ec98f33..5383061db46a 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -3491,8 +3491,9 @@ def test_quantile(): result = pc.quantile(arr, [0.25, 0.5, 0.75], interpolation='linear') assert result.to_pylist() == [1.25, 1.5, 1.75] - with pytest.raises(ValueError, match="Quantile must be between 0 and 1"): - pc.quantile(arr, q=1.1) + for q in (1.1, float('nan')): + with pytest.raises(ValueError, match="Quantile must be between 0 and 1"): + pc.quantile(arr, q=q) with pytest.raises(ValueError, match="not a valid quantile interpolation"): pc.quantile(arr, interpolation='zzz') @@ -3514,6 +3515,14 @@ def test_tdigest(): result = pc.tdigest(arr, [0, 0.5, 1]) # positional `q` assert result.to_pylist() == [1, 2.5, 4] + table = pa.table({'k': [1, 1], 'v': [1.0, 2.0]}) + for q in (1.5, float('nan')): + with pytest.raises(ValueError, match="Quantile must be between 0 and 1"): + pc.tdigest(arr, q=q) + options = pc.TDigestOptions(q=q) + with pytest.raises(ValueError, match="Quantile must be between 0 and 1"): + table.group_by('k').aggregate([('v', 'tdigest', options)]) + def test_fill_null_segfault(): # ARROW-12672