This is an automated email from the ASF dual-hosted git repository.
apitrou pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new bd90043 ARROW-10831: [C++][Compute] Implement quantile kernel
bd90043 is described below
commit bd900437b396a43f66576c61ae7a6b72a9298b66
Author: Yibo Cai <[email protected]>
AuthorDate: Thu Jan 21 15:39:29 2021 +0100
ARROW-10831: [C++][Compute] Implement quantile kernel
Calculate the exact quantile by storing all values and partition
around quantile point at the end. This may require much memory.
As followup task, an approximate method without storing data points
will be implemented.
Closes #8920 from cyb70289/quantile
Lead-authored-by: Yibo Cai <[email protected]>
Co-authored-by: Antoine Pitrou <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
cpp/src/arrow/CMakeLists.txt | 1 +
cpp/src/arrow/compute/api_aggregate.cc | 5 +
cpp/src/arrow/compute/api_aggregate.h | 41 +++
cpp/src/arrow/compute/exec.cc | 23 +-
.../arrow/compute/kernels/aggregate_quantile.cc | 289 +++++++++++++++++++++
cpp/src/arrow/compute/kernels/aggregate_test.cc | 286 ++++++++++++++++++++
cpp/src/arrow/compute/registry.cc | 1 +
cpp/src/arrow/compute/registry_internal.h | 1 +
docs/source/cpp/compute.rst | 8 +-
python/pyarrow/_compute.pyx | 30 +++
python/pyarrow/compute.py | 11 +-
python/pyarrow/includes/libarrow.pxd | 13 +
python/pyarrow/tests/test_compute.py | 39 +++
13 files changed, 727 insertions(+), 21 deletions(-)
diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt
index 7f8b537..eaa6b32 100644
--- a/cpp/src/arrow/CMakeLists.txt
+++ b/cpp/src/arrow/CMakeLists.txt
@@ -365,6 +365,7 @@ if(ARROW_COMPUTE)
compute/registry.cc
compute/kernels/aggregate_basic.cc
compute/kernels/aggregate_mode.cc
+ compute/kernels/aggregate_quantile.cc
compute/kernels/aggregate_var_std.cc
compute/kernels/codegen_internal.cc
compute/kernels/scalar_arithmetic.cc
diff --git a/cpp/src/arrow/compute/api_aggregate.cc
b/cpp/src/arrow/compute/api_aggregate.cc
index 5aeb9f0..586eac2 100644
--- a/cpp/src/arrow/compute/api_aggregate.cc
+++ b/cpp/src/arrow/compute/api_aggregate.cc
@@ -63,5 +63,10 @@ Result<Datum> Variance(const Datum& value, const
VarianceOptions& options,
return CallFunction("variance", {value}, &options, ctx);
}
+Result<Datum> Quantile(const Datum& value, const QuantileOptions& options,
+ ExecContext* ctx) {
+ return CallFunction("quantile", {value}, &options, ctx);
+}
+
} // namespace compute
} // namespace arrow
diff --git a/cpp/src/arrow/compute/api_aggregate.h
b/cpp/src/arrow/compute/api_aggregate.h
index a4cebce..3351861 100644
--- a/cpp/src/arrow/compute/api_aggregate.h
+++ b/cpp/src/arrow/compute/api_aggregate.h
@@ -100,6 +100,33 @@ struct ARROW_EXPORT VarianceOptions : public
FunctionOptions {
int ddof = 0;
};
+/// \brief Control Quantile kernel behavior
+///
+/// By default, returns the median value.
+struct ARROW_EXPORT QuantileOptions : public FunctionOptions {
+ /// Interpolation method to use when quantile lies between two data points
+ enum Interpolation {
+ LINEAR = 0,
+ LOWER,
+ HIGHER,
+ NEAREST,
+ MIDPOINT,
+ };
+
+ explicit QuantileOptions(double q = 0.5, enum Interpolation interpolation =
LINEAR)
+ : q{q}, interpolation{interpolation} {}
+
+ explicit QuantileOptions(std::vector<double> q,
+ enum Interpolation interpolation = LINEAR)
+ : q{std::move(q)}, interpolation{interpolation} {}
+
+ static QuantileOptions Defaults() { return QuantileOptions{}; }
+
+ /// quantile must be between 0 and 1 inclusive
+ std::vector<double> q;
+ enum Interpolation interpolation;
+};
+
/// @}
/// \brief Count non-null (or null) values in an array.
@@ -229,5 +256,19 @@ Result<Datum> Variance(const Datum& value,
const VarianceOptions& options =
VarianceOptions::Defaults(),
ExecContext* ctx = NULLPTR);
+/// \brief Calculate the quantiles of a numeric array
+///
+/// \param[in] value input datum, expecting Array or ChunkedArray
+/// \param[in] options see QuantileOptions for more information
+/// \param[in] ctx the function execution context, optional
+/// \return resulting datum as an array
+///
+/// \since 4.0.0
+/// \note API not yet finalized
+ARROW_EXPORT
+Result<Datum> Quantile(const Datum& value,
+ const QuantileOptions& options =
QuantileOptions::Defaults(),
+ ExecContext* ctx = NULLPTR);
+
} // namespace compute
} // namespace arrow
diff --git a/cpp/src/arrow/compute/exec.cc b/cpp/src/arrow/compute/exec.cc
index fbd8229..ecf3d69 100644
--- a/cpp/src/arrow/compute/exec.cc
+++ b/cpp/src/arrow/compute/exec.cc
@@ -688,10 +688,9 @@ Status PackBatchNoChunks(const std::vector<Datum>& args,
ExecBatch* out) {
switch (arg.kind()) {
case Datum::SCALAR:
case Datum::ARRAY:
+ case Datum::CHUNKED_ARRAY:
length = std::max(arg.length(), length);
break;
- case Datum::CHUNKED_ARRAY:
- return Status::Invalid("Kernel does not support chunked array
arguments");
default:
DCHECK(false);
break;
@@ -722,19 +721,15 @@ class VectorExecutor : public
KernelExecutorImpl<VectorKernel> {
const std::vector<Datum>& outputs) override {
// If execution yielded multiple chunks (because large arrays were split
// based on the ExecContext parameters, then the result is a ChunkedArray
- if (kernel_->output_chunked) {
- if (HaveChunkedArray(inputs) || outputs.size() > 1) {
- return ToChunkedArray(outputs, output_descr_.type);
- } else if (outputs.size() == 1) {
- // Outputs have just one element
- return outputs[0];
- } else {
- // XXX: In the case where no outputs are omitted, is returning a
0-length
- // array always the correct move?
- return MakeArrayOfNull(output_descr_.type, /*length=*/0).ValueOrDie();
- }
- } else {
+ if (kernel_->output_chunked && (HaveChunkedArray(inputs) || outputs.size()
> 1)) {
+ return ToChunkedArray(outputs, output_descr_.type);
+ } else if (outputs.size() == 1) {
+ // Outputs have just one element
return outputs[0];
+ } else {
+ // XXX: In the case where no outputs are omitted, is returning a 0-length
+ // array always the correct move?
+ return MakeArrayOfNull(output_descr_.type, /*length=*/0).ValueOrDie();
}
}
diff --git a/cpp/src/arrow/compute/kernels/aggregate_quantile.cc
b/cpp/src/arrow/compute/kernels/aggregate_quantile.cc
new file mode 100644
index 0000000..98486a8
--- /dev/null
+++ b/cpp/src/arrow/compute/kernels/aggregate_quantile.cc
@@ -0,0 +1,289 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <cmath>
+
+#include "arrow/compute/api_aggregate.h"
+#include "arrow/compute/kernels/common.h"
+#include "arrow/stl_allocator.h"
+#include "arrow/util/bit_run_reader.h"
+
+namespace arrow {
+namespace compute {
+namespace internal {
+
+namespace {
+
+using arrow::internal::checked_pointer_cast;
+using arrow::internal::VisitSetBitRunsVoid;
+
+using QuantileState = internal::OptionsWrapper<QuantileOptions>;
+
+// output is at some input data point, not interpolated
+bool IsDataPoint(const QuantileOptions& options) {
+ // some interpolation methods return exact data point
+ return options.interpolation == QuantileOptions::LOWER ||
+ options.interpolation == QuantileOptions::HIGHER ||
+ options.interpolation == QuantileOptions::NEAREST;
+}
+
+template <typename Dummy, typename InType>
+struct QuantileExecutor {
+ using CType = typename InType::c_type;
+ using Allocator = arrow::stl::allocator<CType>;
+
+ static void Exec(KernelContext* ctx, const ExecBatch& batch, Datum* out) {
+ // validate arguments
+ if (ctx->state() == nullptr) {
+ ctx->SetStatus(Status::Invalid("Quantile requires QuantileOptions"));
+ return;
+ }
+
+ const QuantileOptions& options = QuantileState::Get(ctx);
+ if (options.q.empty()) {
+ ctx->SetStatus(Status::Invalid("Requires quantile argument"));
+ return;
+ }
+ for (double q : options.q) {
+ if (q < 0 || q > 1) {
+ ctx->SetStatus(Status::Invalid("Quantile must be between 0 and 1"));
+ return;
+ }
+ }
+
+ // copy all chunks to a buffer, ignore nulls and nans
+ std::vector<CType, Allocator> in_buffer(Allocator(ctx->memory_pool()));
+
+ const Datum& datum = batch[0];
+ const int64_t in_length = datum.length() - datum.null_count();
+ if (in_length > 0) {
+ in_buffer.resize(in_length);
+
+ int64_t index = 0;
+ for (const auto& array : datum.chunks()) {
+ index += CopyArray(in_buffer.data() + index, *array);
+ }
+ DCHECK_EQ(index, in_length);
+
+ // drop nan
+ if (is_floating_type<InType>::value) {
+ const auto& it = std::remove_if(in_buffer.begin(), in_buffer.end(),
+ [](CType v) { return v != v; });
+ in_buffer.resize(it - in_buffer.begin());
+ }
+ }
+
+ // prepare out array
+ int64_t out_length = options.q.size();
+ if (in_buffer.empty()) {
+ out_length = 0; // input is empty or only contains null and nan, return
empty array
+ }
+ // out type depends on options
+ const bool is_datapoint = IsDataPoint(options);
+ std::shared_ptr<DataType> out_type;
+ if (is_datapoint) {
+ out_type = TypeTraits<InType>::type_singleton();
+ } else {
+ out_type = float64();
+ }
+ auto out_data = ArrayData::Make(out_type, out_length, 0);
+ out_data->buffers.resize(2, nullptr);
+
+ // calculate quantiles
+ if (out_length > 0) {
+ const auto out_bit_width =
checked_pointer_cast<NumberType>(out_type)->bit_width();
+ KERNEL_ASSIGN_OR_RAISE(out_data->buffers[1], ctx,
+ ctx->Allocate(out_length * out_bit_width / 8));
+
+ // find quantiles in descending order
+ std::vector<int64_t> q_indices(out_length);
+ std::iota(q_indices.begin(), q_indices.end(), 0);
+ std::sort(q_indices.begin(), q_indices.end(),
+ [&options](int64_t left_index, int64_t right_index) {
+ return options.q[right_index] < options.q[left_index];
+ });
+
+ // input array is partitioned around data point at `last_index` (pivot)
+ // for next quatile which is smaller, we only consider inputs left of
the pivot
+ uint64_t last_index = in_buffer.size();
+ if (is_datapoint) {
+ CType* out_buffer = out_data->template GetMutableValues<CType>(1);
+ for (int64_t i = 0; i < out_length; ++i) {
+ const int64_t q_index = q_indices[i];
+ out_buffer[q_index] = GetQuantileAtDataPoint(
+ in_buffer, &last_index, options.q[q_index],
options.interpolation);
+ }
+ } else {
+ double* out_buffer = out_data->template GetMutableValues<double>(1);
+ for (int64_t i = 0; i < out_length; ++i) {
+ const int64_t q_index = q_indices[i];
+ out_buffer[q_index] = GetQuantileByInterp(
+ in_buffer, &last_index, options.q[q_index],
options.interpolation);
+ }
+ }
+ }
+
+ *out = Datum(std::move(out_data));
+ }
+
+ static int64_t CopyArray(CType* buffer, const Array& array) {
+ const int64_t n = array.length() - array.null_count();
+ if (n > 0) {
+ int64_t index = 0;
+ const ArrayData& data = *array.data();
+ const CType* values = data.GetValues<CType>(1);
+ VisitSetBitRunsVoid(data.buffers[0], data.offset, data.length,
+ [&](int64_t pos, int64_t len) {
+ memcpy(buffer + index, values + pos, len *
sizeof(CType));
+ index += len;
+ });
+ DCHECK_EQ(index, n);
+ }
+ return n;
+ }
+
+ // return quantile located exactly at some input data point
+ static CType GetQuantileAtDataPoint(std::vector<CType, Allocator>& in,
+ uint64_t* last_index, double q,
+ enum QuantileOptions::Interpolation
interpolation) {
+ const double index = (in.size() - 1) * q;
+ uint64_t datapoint_index = static_cast<uint64_t>(index);
+ const double fraction = index - datapoint_index;
+
+ if (interpolation == QuantileOptions::LINEAR ||
+ interpolation == QuantileOptions::MIDPOINT) {
+ DCHECK_EQ(fraction, 0);
+ }
+
+ // convert NEAREST interpolation method to LOWER or HIGHER
+ if (interpolation == QuantileOptions::NEAREST) {
+ if (fraction < 0.5) {
+ interpolation = QuantileOptions::LOWER;
+ } else if (fraction > 0.5) {
+ interpolation = QuantileOptions::HIGHER;
+ } else {
+ // round 0.5 to nearest even number, similar to numpy.around
+ interpolation =
+ (datapoint_index & 1) ? QuantileOptions::HIGHER :
QuantileOptions::LOWER;
+ }
+ }
+
+ if (interpolation == QuantileOptions::HIGHER && fraction != 0) {
+ ++datapoint_index;
+ }
+
+ if (datapoint_index != *last_index) {
+ DCHECK_LT(datapoint_index, *last_index);
+ std::nth_element(in.begin(), in.begin() + datapoint_index,
+ in.begin() + *last_index);
+ *last_index = datapoint_index;
+ }
+
+ return in[datapoint_index];
+ }
+
+ // return quantile interpolated from adjacent input data points
+ static double GetQuantileByInterp(std::vector<CType, Allocator>& in,
+ uint64_t* last_index, double q,
+ enum QuantileOptions::Interpolation
interpolation) {
+ const double index = (in.size() - 1) * q;
+ const uint64_t lower_index = static_cast<uint64_t>(index);
+ const double fraction = index - lower_index;
+
+ if (lower_index != *last_index) {
+ DCHECK_LT(lower_index, *last_index);
+ std::nth_element(in.begin(), in.begin() + lower_index, in.begin() +
*last_index);
+ }
+
+ const double lower_value = static_cast<double>(in[lower_index]);
+ if (fraction == 0) {
+ *last_index = lower_index;
+ return lower_value;
+ }
+
+ const uint64_t higher_index = lower_index + 1;
+ DCHECK_LT(higher_index, in.size());
+ if (lower_index != *last_index && higher_index != *last_index) {
+ DCHECK_LT(higher_index, *last_index);
+ // higher value must be the minimal value after lower_index
+ auto min = std::min_element(in.begin() + higher_index, in.begin() +
*last_index);
+ std::iter_swap(in.begin() + higher_index, min);
+ }
+ *last_index = lower_index;
+
+ const double higher_value = static_cast<double>(in[higher_index]);
+
+ if (interpolation == QuantileOptions::LINEAR) {
+ // more stable than naive linear interpolation
+ return fraction * higher_value + (1 - fraction) * lower_value;
+ } else if (interpolation == QuantileOptions::MIDPOINT) {
+ return lower_value / 2 + higher_value / 2;
+ } else {
+ DCHECK(false);
+ return NAN;
+ }
+ }
+};
+
+Result<ValueDescr> ResolveOutput(KernelContext* ctx,
+ const std::vector<ValueDescr>& args) {
+ const QuantileOptions& options = QuantileState::Get(ctx);
+ if (IsDataPoint(options)) {
+ return ValueDescr::Array(args[0].type);
+ } else {
+ return ValueDescr::Array(float64());
+ }
+}
+
+void AddQuantileKernels(VectorFunction* func) {
+ VectorKernel base;
+ base.init = QuantileState::Init;
+ base.can_execute_chunkwise = false;
+ base.output_chunked = false;
+
+ for (const auto& ty : NumericTypes()) {
+ base.signature =
+ KernelSignature::Make({InputType::Array(ty)},
OutputType(ResolveOutput));
+ // output type is determined at runtime, set template argument to nulltype
+ base.exec = GenerateNumeric<QuantileExecutor, NullType>(*ty);
+ DCHECK_OK(func->AddKernel(base));
+ }
+}
+
+const FunctionDoc quantile_doc{
+ "Compute an array of quantiles of a numeric array or chunked array",
+ ("By default, 0.5 quantile (median) is returned.\n"
+ "If quantile lies between two data points, an interpolated value is\n"
+ "returned based on selected interpolation method.\n"
+ "Nulls and NaNs are ignored.\n"
+ "An empty array is returned if there is no valid data point."),
+ {"array"},
+ "QuantileOptions"};
+
+} // namespace
+
+void RegisterScalarAggregateQuantile(FunctionRegistry* registry) {
+ static QuantileOptions default_options;
+ auto func = std::make_shared<VectorFunction>("quantile", Arity::Unary(),
&quantile_doc,
+ &default_options);
+ AddQuantileKernels(func.get());
+ DCHECK_OK(registry->AddFunction(std::move(func)));
+}
+
+} // namespace internal
+} // namespace compute
+} // namespace arrow
diff --git a/cpp/src/arrow/compute/kernels/aggregate_test.cc
b/cpp/src/arrow/compute/kernels/aggregate_test.cc
index 57b2931..e944166 100644
--- a/cpp/src/arrow/compute/kernels/aggregate_test.cc
+++ b/cpp/src/arrow/compute/kernels/aggregate_test.cc
@@ -1321,5 +1321,291 @@ TEST_F(TestVarStdKernelIntegerLength, Basics) {
}
#endif
+//
+// Quantile
+//
+
+template <typename ArrowType>
+class TestPrimitiveQuantileKernel : public ::testing::Test {
+ public:
+ using Traits = TypeTraits<ArrowType>;
+ using CType = typename ArrowType::c_type;
+
+ void AssertQuantilesAre(const Datum& array, QuantileOptions options,
+ const std::vector<std::vector<Datum>>& expected) {
+ ASSERT_EQ(options.q.size(), expected.size());
+
+ for (size_t i = 0; i < this->interpolations_.size(); ++i) {
+ options.interpolation = this->interpolations_[i];
+
+ ASSERT_OK_AND_ASSIGN(Datum out, Quantile(array, options));
+ const auto& out_array = out.make_array();
+ ASSERT_OK(out_array->ValidateFull());
+ ASSERT_EQ(out_array->length(), options.q.size());
+ ASSERT_EQ(out_array->null_count(), 0);
+ ASSERT_EQ(out_array->type(), expected[0][i].type());
+
+ if (out_array->type() == float64()) {
+ const double* quantiles = out_array->data()->GetValues<double>(1);
+ for (int64_t j = 0; j < out_array->length(); ++j) {
+ const auto& numeric_scalar =
+ std::static_pointer_cast<DoubleScalar>(expected[j][i].scalar());
+ ASSERT_TRUE((quantiles[j] == numeric_scalar->value) ||
+ (std::isnan(quantiles[j]) &&
std::isnan(numeric_scalar->value)));
+ }
+ } else {
+ ASSERT_EQ(out_array->type(), type_singleton());
+ const CType* quantiles = out_array->data()->GetValues<CType>(1);
+ for (int64_t j = 0; j < out_array->length(); ++j) {
+ const auto& numeric_scalar =
+
std::static_pointer_cast<NumericScalar<ArrowType>>(expected[j][i].scalar());
+ ASSERT_EQ(quantiles[j], numeric_scalar->value);
+ }
+ }
+ }
+ }
+
+ void AssertQuantilesAre(const std::string& json, const std::vector<double>&
q,
+ const std::vector<std::vector<Datum>>& expected) {
+ auto array = ArrayFromJSON(type_singleton(), json);
+ AssertQuantilesAre(array, QuantileOptions{q}, expected);
+ }
+
+ void AssertQuantilesAre(const std::vector<std::string>& json,
+ const std::vector<double>& q,
+ const std::vector<std::vector<Datum>>& expected) {
+ auto chunked = ChunkedArrayFromJSON(type_singleton(), json);
+ AssertQuantilesAre(chunked, QuantileOptions{q}, expected);
+ }
+
+ void AssertQuantileIs(const Datum& array, double q,
+ const std::vector<Datum>& expected) {
+ AssertQuantilesAre(array, QuantileOptions{q}, {expected});
+ }
+
+ void AssertQuantileIs(const std::string& json, double q,
+ const std::vector<Datum>& expected) {
+ auto array = ArrayFromJSON(type_singleton(), json);
+ AssertQuantileIs(array, q, expected);
+ }
+
+ void AssertQuantileIs(const std::vector<std::string>& json, double q,
+ const std::vector<Datum>& expected) {
+ auto chunked = ChunkedArrayFromJSON(type_singleton(), json);
+ AssertQuantileIs(chunked, q, expected);
+ }
+
+ void AssertQuantilesEmpty(const Datum& array, const std::vector<double>& q) {
+ QuantileOptions options{q};
+ for (auto interpolation : this->interpolations_) {
+ options.interpolation = interpolation;
+ ASSERT_OK_AND_ASSIGN(Datum out, Quantile(array, options));
+ ASSERT_OK(out.make_array()->ValidateFull());
+ ASSERT_EQ(out.array()->length, 0);
+ }
+ }
+
+ void AssertQuantilesEmpty(const std::string& json, const
std::vector<double>& q) {
+ auto array = ArrayFromJSON(type_singleton(), json);
+ AssertQuantilesEmpty(array, q);
+ }
+
+ void AssertQuantilesEmpty(const std::vector<std::string>& json,
+ const std::vector<double>& q) {
+ auto chunked = ChunkedArrayFromJSON(type_singleton(), json);
+ AssertQuantilesEmpty(chunked, q);
+ }
+
+ std::shared_ptr<DataType> type_singleton() { return
Traits::type_singleton(); }
+
+ std::vector<enum QuantileOptions::Interpolation> interpolations_{
+ QuantileOptions::LINEAR, QuantileOptions::LOWER, QuantileOptions::HIGHER,
+ QuantileOptions::NEAREST, QuantileOptions::MIDPOINT};
+};
+
+template <typename ArrowType>
+class TestIntegerQuantileKernel : public
TestPrimitiveQuantileKernel<ArrowType> {};
+
+template <typename ArrowType>
+class TestFloatingQuantileKernel : public
TestPrimitiveQuantileKernel<ArrowType> {};
+
+template <typename ArrowType>
+class TestInt64QuantileKernel : public TestPrimitiveQuantileKernel<ArrowType>
{};
+
+#define INTYPE(x) Datum(static_cast<typename TypeParam::c_type>(x))
+#define DOUBLE(x) Datum(static_cast<double>(x))
+// output type per interplation: linear, lower, higher, nearest, midpoint
+#define O(a, b, c, d, e) \
+ { DOUBLE(a), INTYPE(b), INTYPE(c), INTYPE(d), DOUBLE(e) }
+
+TYPED_TEST_SUITE(TestIntegerQuantileKernel, IntegralArrowTypes);
+TYPED_TEST(TestIntegerQuantileKernel, Basics) {
+ // reference values from numpy
+ // ordered by interpolation method: {linear, lower, higher, nearest,
midpoint}
+ this->AssertQuantileIs("[1]", 0.1, O(1, 1, 1, 1, 1));
+ this->AssertQuantileIs("[1, 2]", 0.5, O(1.5, 1, 2, 1, 1.5));
+ this->AssertQuantileIs("[3, 5, 2, 9, 0, 1, 8]", 0.5, O(3, 3, 3, 3, 3));
+ this->AssertQuantileIs("[3, 5, 2, 9, 0, 1, 8]", 0.33, O(1.98, 1, 2, 2, 1.5));
+ this->AssertQuantileIs("[3, 5, 2, 9, 0, 1, 8]", 0.9, O(8.4, 8, 9, 8, 8.5));
+ this->AssertQuantilesAre("[3, 5, 2, 9, 0, 1, 8]", {0.5, 0.9},
+ {O(3, 3, 3, 3, 3), O(8.4, 8, 9, 8, 8.5)});
+ this->AssertQuantilesAre("[3, 5, 2, 9, 0, 1, 8]", {1, 0.5},
+ {O(9, 9, 9, 9, 9), O(3, 3, 3, 3, 3)});
+ this->AssertQuantileIs("[3, 5, 2, 9, 0, 1, 8]", 0, O(0, 0, 0, 0, 0));
+ this->AssertQuantileIs("[3, 5, 2, 9, 0, 1, 8]", 1, O(9, 9, 9, 9, 9));
+
+ this->AssertQuantileIs("[5, null, null, 3, 9, null, 8, 1, 2, 0]", 0.21,
+ O(1.26, 1, 2, 1, 1.5));
+ this->AssertQuantilesAre("[5, null, null, 3, 9, null, 8, 1, 2, 0]", {0.5,
0.9},
+ {O(3, 3, 3, 3, 3), O(8.4, 8, 9, 8, 8.5)});
+ this->AssertQuantilesAre("[5, null, null, 3, 9, null, 8, 1, 2, 0]", {0.9,
0.5},
+ {O(8.4, 8, 9, 8, 8.5), O(3, 3, 3, 3, 3)});
+
+ this->AssertQuantileIs({"[5]", "[null, null]", "[3, 9, null]", "[8, 1, 2,
0]"}, 0.33,
+ O(1.98, 1, 2, 2, 1.5));
+ this->AssertQuantilesAre({"[5]", "[null, null]", "[3, 9, null]", "[8, 1, 2,
0]"},
+ {0.21, 1}, {O(1.26, 1, 2, 1, 1.5), O(9, 9, 9, 9,
9)});
+
+ this->AssertQuantilesEmpty("[]", {0.5});
+ this->AssertQuantilesEmpty("[null, null, null]", {0.1, 0.2});
+ this->AssertQuantilesEmpty({"[null, null]", "[]", "[null]"}, {0.3, 0.4});
+}
+
+#ifndef __MINGW32__
+TYPED_TEST_SUITE(TestFloatingQuantileKernel, RealArrowTypes);
+TYPED_TEST(TestFloatingQuantileKernel, Floats) {
+ // ordered by interpolation method: {linear, lower, higher, nearest,
midpoint}
+ this->AssertQuantileIs("[-9, 7, Inf, -Inf, 2, 11]", 0.5, O(4.5, 2, 7, 2,
4.5));
+ this->AssertQuantileIs("[-9, 7, Inf, -Inf, 2, 11]", 0.1,
+ O(-INFINITY, -INFINITY, -9, -INFINITY, -INFINITY));
+ this->AssertQuantileIs("[-9, 7, Inf, -Inf, 2, 11]", 0.9,
+ O(INFINITY, 11, INFINITY, 11, INFINITY));
+ this->AssertQuantilesAre("[-9, 7, Inf, -Inf, 2, 11]", {0.3, 0.6},
+ {O(-3.5, -9, 2, 2, -3.5), O(7, 7, 7, 7, 7)});
+ this->AssertQuantileIs("[-Inf, Inf]", 0.2, O(NAN, -INFINITY, INFINITY,
-INFINITY, NAN));
+
+ this->AssertQuantileIs("[NaN, -9, 7, Inf, null, null, -Inf, NaN, 2, 11]",
0.5,
+ O(4.5, 2, 7, 2, 4.5));
+ this->AssertQuantilesAre("[null, -9, 7, Inf, NaN, NaN, -Inf, null, 2, 11]",
{0.3, 0.6},
+ {O(-3.5, -9, 2, 2, -3.5), O(7, 7, 7, 7, 7)});
+ this->AssertQuantilesAre("[null, -9, 7, Inf, NaN, NaN, -Inf, null, 2, 11]",
{0.6, 0.3},
+ {O(7, 7, 7, 7, 7), O(-3.5, -9, 2, 2, -3.5)});
+
+ this->AssertQuantileIs({"[NaN, -9, 7, Inf]", "[null, NaN]", "[-Inf, NaN, 2,
11]"}, 0.5,
+ O(4.5, 2, 7, 2, 4.5));
+ this->AssertQuantilesAre({"[null, -9, 7, Inf]", "[NaN, NaN]", "[-Inf, null,
2, 11]"},
+ {0.3, 0.6}, {O(-3.5, -9, 2, 2, -3.5), O(7, 7, 7, 7,
7)});
+
+ this->AssertQuantilesEmpty("[]", {0.5, 0.6});
+ this->AssertQuantilesEmpty("[null, NaN, null]", {0.1});
+ this->AssertQuantilesEmpty({"[NaN, NaN]", "[]", "[null]"}, {0.3, 0.4});
+}
+#endif
+
+// Test big int64 numbers cannot be precisely presented by double
+TYPED_TEST_SUITE(TestInt64QuantileKernel, Int64Type);
+TYPED_TEST(TestInt64QuantileKernel, Int64) {
+ this->AssertQuantileIs(
+ "[9223372036854775806, 9223372036854775807]", 0.5,
+ O(9.223372036854776e+18, 9223372036854775806, 9223372036854775807,
+ 9223372036854775806, 9.223372036854776e+18));
+}
+
+#undef INTYPE
+#undef DOUBLE
+#undef O
+
+#ifndef __MINGW32__
+class TestRandomQuantileKernel : public TestPrimitiveQuantileKernel<Int32Type>
{
+ public:
+ void CheckQuantiles(int64_t array_size, int64_t num_quantiles) {
+ auto rand = random::RandomArrayGenerator(0x5487658);
+ // set a small value range to exercise input array with equal values
+ const auto array = rand.Numeric<Int32Type>(array_size, -100, 200, 0.1);
+
+ std::vector<double> quantiles;
+ random_real(num_quantiles, 0x5487658, 0.0, 1.0, &quantiles);
+ // make sure to exercise 0 and 1 quantiles
+ *std::min_element(quantiles.begin(), quantiles.end()) = 0;
+ *std::max_element(quantiles.begin(), quantiles.end()) = 1;
+
+ this->AssertQuantilesAre(array, QuantileOptions{quantiles},
+ NaiveQuantile(*array, quantiles));
+ }
+
+ private:
+ std::vector<std::vector<Datum>> NaiveQuantile(const Array& array,
+ const std::vector<double>&
quantiles) {
+ // copy and sort input array
+ std::vector<int32_t> input(array.length() - array.null_count());
+ const int32_t* values = array.data()->GetValues<int32_t>(1);
+ const auto bitmap = array.null_bitmap_data();
+ int64_t index = 0;
+ for (int64_t i = 0; i < array.length(); ++i) {
+ if (BitUtil::GetBit(bitmap, i)) {
+ input[index++] = values[i];
+ }
+ }
+ std::sort(input.begin(), input.end());
+
+ std::vector<std::vector<Datum>> output(quantiles.size(),
+
std::vector<Datum>(interpolations_.size()));
+ for (uint64_t i = 0; i < interpolations_.size(); ++i) {
+ const auto interp = interpolations_[i];
+ for (uint64_t j = 0; j < quantiles.size(); ++j) {
+ output[j][i] = GetQuantile(input, quantiles[j], interp);
+ }
+ }
+ return output;
+ }
+
+ Datum GetQuantile(const std::vector<int32_t>& input, double q,
+ enum QuantileOptions::Interpolation interp) {
+ const double index = (input.size() - 1) * q;
+ const uint64_t lower_index = static_cast<uint64_t>(index);
+ const double fraction = index - lower_index;
+
+ switch (interp) {
+ case QuantileOptions::LOWER:
+ return Datum(input[lower_index]);
+ case QuantileOptions::HIGHER:
+ return Datum(input[lower_index + (fraction != 0)]);
+ case QuantileOptions::NEAREST:
+ if (fraction < 0.5) {
+ return Datum(input[lower_index]);
+ } else if (fraction > 0.5) {
+ return Datum(input[lower_index + 1]);
+ } else {
+ return Datum(input[lower_index + (lower_index & 1)]);
+ }
+ case QuantileOptions::LINEAR:
+ if (fraction == 0) {
+ return Datum(static_cast<double>(input[lower_index]));
+ } else {
+ return Datum(fraction * input[lower_index + 1] +
+ (1 - fraction) * input[lower_index]);
+ }
+ case QuantileOptions::MIDPOINT:
+ if (fraction == 0) {
+ return Datum(static_cast<double>(input[lower_index]));
+ } else {
+ return Datum(input[lower_index] / 2.0 + input[lower_index + 1] /
2.0);
+ }
+ default:
+ return Datum(NAN);
+ }
+ }
+};
+
+TEST_F(TestRandomQuantileKernel, Normal) {
+ this->CheckQuantiles(/*array_size=*/10000, /*num_quantiles=*/100);
+}
+
+TEST_F(TestRandomQuantileKernel, Overlapped) {
+ // much more quantiles than array size => many overlaps
+ this->CheckQuantiles(/*array_size=*/999, /*num_quantiles=*/9999);
+}
+#endif
+
} // namespace compute
} // namespace arrow
diff --git a/cpp/src/arrow/compute/registry.cc
b/cpp/src/arrow/compute/registry.cc
index 7ef1e26..b1e0d48 100644
--- a/cpp/src/arrow/compute/registry.cc
+++ b/cpp/src/arrow/compute/registry.cc
@@ -129,6 +129,7 @@ static std::unique_ptr<FunctionRegistry>
CreateBuiltInRegistry() {
// Aggregate functions
RegisterScalarAggregateBasic(registry.get());
RegisterScalarAggregateMode(registry.get());
+ RegisterScalarAggregateQuantile(registry.get());
RegisterScalarAggregateVariance(registry.get());
// Vector functions
diff --git a/cpp/src/arrow/compute/registry_internal.h
b/cpp/src/arrow/compute/registry_internal.h
index 78e134e..4e39eeb 100644
--- a/cpp/src/arrow/compute/registry_internal.h
+++ b/cpp/src/arrow/compute/registry_internal.h
@@ -44,6 +44,7 @@ void RegisterVectorSort(FunctionRegistry* registry);
// Aggregate functions
void RegisterScalarAggregateBasic(FunctionRegistry* registry);
void RegisterScalarAggregateMode(FunctionRegistry* registry);
+void RegisterScalarAggregateQuantile(FunctionRegistry* registry);
void RegisterScalarAggregateVariance(FunctionRegistry* registry);
} // namespace internal
diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst
index 09e7737..c513ed5 100644
--- a/docs/source/cpp/compute.rst
+++ b/docs/source/cpp/compute.rst
@@ -150,9 +150,11 @@ Aggregations
+--------------------------+------------+--------------------+-----------------------+--------------------------------------------+
| mode | Unary | Numeric | Struct (2)
| :struct:`ModeOptions` |
+--------------------------+------------+--------------------+-----------------------+--------------------------------------------+
+| quantile | Unary | Numeric | Scalar Numeric
(3) | :struct:`QuantileOptions` |
++--------------------------+------------+--------------------+-----------------------+--------------------------------------------+
| stddev | Unary | Numeric | Scalar Float64
| :struct:`VarianceOptions` |
+--------------------------+------------+--------------------+-----------------------+--------------------------------------------+
-| sum | Unary | Numeric | Scalar Numeric
(3) | |
+| sum | Unary | Numeric | Scalar Numeric
(4) | |
+--------------------------+------------+--------------------+-----------------------+--------------------------------------------+
| variance | Unary | Numeric | Scalar Float64
| :struct:`VarianceOptions` |
+--------------------------+------------+--------------------+-----------------------+--------------------------------------------+
@@ -168,7 +170,9 @@ Notes:
Note that the output can have less than *N* elements if the input has
less than *N* distinct values.
-* \(3) Output is Int64, UInt64 or Float64, depending on the input type
+* \(3) Output is Float64 or input type, depending on QuantileOptions.
+
+* \(4) Output is Int64, UInt64 or Float64, depending on the input type.
Element-wise ("scalar") functions
---------------------------------
diff --git a/python/pyarrow/_compute.pyx b/python/pyarrow/_compute.pyx
index ecc82c2..90b5eeb 100644
--- a/python/pyarrow/_compute.pyx
+++ b/python/pyarrow/_compute.pyx
@@ -930,3 +930,33 @@ class SortOptions(_SortOptions):
if sort_keys is None:
sort_keys = []
self._set_options(sort_keys)
+
+
+cdef class _QuantileOptions(FunctionOptions):
+ cdef:
+ CQuantileOptions quantile_options
+
+ cdef const CFunctionOptions* get_options(self) except NULL:
+ return &self.quantile_options
+
+ def _set_options(self, quantiles, interp):
+ interp_dict = {
+ 'linear': CQuantileInterp_LINEAR,
+ 'lower': CQuantileInterp_LOWER,
+ 'higher': CQuantileInterp_HIGHER,
+ 'nearest': CQuantileInterp_NEAREST,
+ 'midpoint': CQuantileInterp_MIDPOINT,
+ }
+ if interp not in interp_dict:
+ raise ValueError(
+ '{!r} is not a valid interpolation'
+ .format(interp))
+ self.quantile_options.interpolation = interp_dict[interp]
+ self.quantile_options.q = quantiles
+
+
+class QuantileOptions(_QuantileOptions):
+ def __init__(self, *, q=0.5, interpolation='linear'):
+ if not isinstance(q, (list, tuple, np.ndarray)):
+ q = [q]
+ self._set_options(q, interpolation)
diff --git a/python/pyarrow/compute.py b/python/pyarrow/compute.py
index b09be7d..e1e64a6 100644
--- a/python/pyarrow/compute.py
+++ b/python/pyarrow/compute.py
@@ -27,23 +27,24 @@ from pyarrow._compute import ( # noqa
VectorFunction,
VectorKernel,
# Option classes
+ ArraySortOptions,
CastOptions,
CountOptions,
FilterOptions,
MatchSubstringOptions,
- SplitOptions,
- SplitPatternOptions,
- TrimOptions,
MinMaxOptions,
ModeOptions,
+ SplitOptions,
+ SplitPatternOptions,
PartitionNthOptions,
ProjectOptions,
+ QuantileOptions,
SetLookupOptions,
+ SortOptions,
StrptimeOptions,
TakeOptions,
+ TrimOptions,
VarianceOptions,
- ArraySortOptions,
- SortOptions,
# Functions
function_registry,
call_function,
diff --git a/python/pyarrow/includes/libarrow.pxd
b/python/pyarrow/includes/libarrow.pxd
index 976af53..b1372d0 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -1848,6 +1848,19 @@ cdef extern from "arrow/compute/api.h" namespace
"arrow::compute" nogil:
"arrow::compute::SortOptions"(CFunctionOptions):
vector[CSortKey] sort_keys
+ enum CQuantileInterp \
+ "arrow::compute::QuantileOptions::Interpolation":
+ CQuantileInterp_LINEAR "arrow::compute::QuantileOptions::LINEAR"
+ CQuantileInterp_LOWER "arrow::compute::QuantileOptions::LOWER"
+ CQuantileInterp_HIGHER "arrow::compute::QuantileOptions::HIGHER"
+ CQuantileInterp_NEAREST "arrow::compute::QuantileOptions::NEAREST"
+ CQuantileInterp_MIDPOINT "arrow::compute::QuantileOptions::MIDPOINT"
+
+ cdef cppclass CQuantileOptions \
+ "arrow::compute::QuantileOptions"(CFunctionOptions):
+ vector[double] q
+ CQuantileInterp interpolation
+
enum DatumType" arrow::Datum::type":
DatumType_NONE" arrow::Datum::NONE"
DatumType_SCALAR" arrow::Datum::SCALAR"
diff --git a/python/pyarrow/tests/test_compute.py
b/python/pyarrow/tests/test_compute.py
index b58844b..06a0269 100644
--- a/python/pyarrow/tests/test_compute.py
+++ b/python/pyarrow/tests/test_compute.py
@@ -1160,3 +1160,42 @@ def test_index_in():
result = pc.index_in(arr, value_set=pa.array([1, 3]), skip_nulls=True)
assert result.to_pylist() == [0, None, None, 0, None, 1]
+
+
+def test_quantile():
+ arr = pa.array([1, 2, 3, 4])
+
+ result = pc.quantile(arr)
+ assert result.to_pylist() == [2.5]
+
+ result = pc.quantile(arr, interpolation='lower')
+ assert result.to_pylist() == [2]
+ result = pc.quantile(arr, interpolation='higher')
+ assert result.to_pylist() == [3]
+ result = pc.quantile(arr, interpolation='nearest')
+ assert result.to_pylist() == [3]
+ result = pc.quantile(arr, interpolation='midpoint')
+ assert result.to_pylist() == [2.5]
+ result = pc.quantile(arr, interpolation='linear')
+ assert result.to_pylist() == [2.5]
+
+ arr = pa.array([1, 2])
+
+ result = pc.quantile(arr, q=[0.25, 0.5, 0.75])
+ assert result.to_pylist() == [1.25, 1.5, 1.75]
+
+ result = pc.quantile(arr, q=[0.25, 0.5, 0.75], interpolation='lower')
+ assert result.to_pylist() == [1, 1, 1]
+ result = pc.quantile(arr, q=[0.25, 0.5, 0.75], interpolation='higher')
+ assert result.to_pylist() == [2, 2, 2]
+ result = pc.quantile(arr, q=[0.25, 0.5, 0.75], interpolation='midpoint')
+ assert result.to_pylist() == [1.5, 1.5, 1.5]
+ result = pc.quantile(arr, q=[0.25, 0.5, 0.75], interpolation='nearest')
+ assert result.to_pylist() == [1, 1, 2]
+ result = pc.quantile(arr, q=[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)
+ with pytest.raises(ValueError, match="'zzz' is not a valid interpolation"):
+ pc.quantile(arr, interpolation='zzz')
