cyb70289 commented on a change in pull request #8920:
URL: https://github.com/apache/arrow/pull/8920#discussion_r555822507
##########
File path: cpp/src/arrow/compute/kernels/aggregate_test.cc
##########
@@ -1321,5 +1321,288 @@ 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() == 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);
+ }
+ } else {
+ ASSERT_EQ(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_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) }
+// output type same as input if only 0 and 1 quantiles are calculated
+#define I(a, b, c, d, e) \
+ { INTYPE(a), INTYPE(b), INTYPE(c), INTYPE(d), INTYPE(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, I(0, 0, 0, 0, 0));
+ this->AssertQuantileIs("[3, 5, 2, 9, 0, 1, 8]", 1, I(9, 9, 9, 9, 9));
+ this->AssertQuantilesAre("[3, 5, 2, 9, 0, 1, 8]", {1, 0},
+ {I(9, 9, 9, 9, 9), I(0, 0, 0, 0, 0)});
+ this->AssertQuantilesAre(
+ "[3, 5, 2, 9, 0, 1, 8]", {1, 0, 0, 1},
+ {I(9, 9, 9, 9, 9), I(0, 0, 0, 0, 0), I(0, 0, 0, 0, 0), I(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) {
+ 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("[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});
+}
Review comment:
Yes, will change.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
