GH-55372: [C++][Python] Reject q outside [0, 1] in tdigest and hash_tdigest

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<double>::quiet_NaN()}) {
+    auto options = std::make_shared<TDigestOptions>(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)),

cpp/src/arrow/compute/kernels/aggregate_internal.h

@@ -19,8 +19,10 @@
 
 #include <cmath>
 #include <initializer_list>
+#include <vector>
 
 #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<double>& 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 <typename I, typename Enable = void>
 struct FindAccumulatorType {};

cpp/src/arrow/compute/kernels/aggregate_quantile.cc

@@ -20,6 +20,7 @@
 #include <vector>
 
 #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 <typename InType>

cpp/src/arrow/compute/kernels/aggregate_tdigest.cc

@@ -167,8 +167,9 @@ struct TDigestInitState {
 
 Result<std::unique_ptr<KernelState>> TDigestInit(KernelContext* ctx,
                                                  const KernelInitArgs& args) {
-  TDigestInitState visitor(ctx, *args.inputs[0].type,
-                           static_cast<const TDigestOptions&>(*args.options));
+  const auto& options = static_cast<const TDigestOptions&>(*args.options);
+  RETURN_NOT_OK(ValidateQuantileLevels(options.q));
+  TDigestInitState visitor(ctx, *args.inputs[0].type, options);
   return visitor.Create();
 }
 

cpp/src/arrow/compute/kernels/aggregate_test.cc

@@ -4399,6 +4399,20 @@ TEST(TestQuantileKernel, Scalar) {
   }
 }
 
+TEST(TestQuantileKernel, InvalidQuantileLevel) {
+  const double nan = std::numeric_limits<double>::quiet_NaN();
+  auto values = ArrayFromJSON(float64(), "[1.0, 2.0, 3.0]");
+  auto empty = ArrayFromJSON(float64(), "[]");
+  for (const std::vector<double>& q :
+       std::vector<std::vector<double>>{{-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<double>::quiet_NaN();
+  auto values = ArrayFromJSON(float64(), "[1.0, 2.0, 3.0]");
+  auto empty = ArrayFromJSON(float64(), "[]");
+  for (const std::vector<double>& q :
+       std::vector<std::vector<double>>{{-0.1},
+                                        {1.1},
+                                        {0.5, 1.5},
+                                        {-std::numeric_limits<double>::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<double>{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()}) {

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<const TDigestOptions*>(args.options);
+    RETURN_NOT_OK(ValidateQuantileLevels(options_.q));
     if (is_decimal_type<Type>::value) {
       decimal_scale_ = checked_cast<const DecimalType&>(*args.inputs[0].type).scale();
     } else {

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