Taepper commented on code in PR #46926:
URL: https://github.com/apache/arrow/pull/46926#discussion_r3404257475
##########
cpp/src/arrow/compute/kernels/select_k_test.cc:
##########
@@ -283,6 +380,62 @@ TYPED_TEST(TestSelectKWithChunkedArray,
RandomValuesWithSlices) {
}
}
+TYPED_TEST(TestSelectKWithChunkedArray, PartialSelectKNull) {
+ auto chunked_array = ChunkedArrayFromJSON(uint8(), {
+ "[null, 1]",
+ "[3, null, 2]",
+ "[1]",
+ });
+ std::vector<SortKey> sort_keys{SortKey("a", SortOrder::Ascending)};
+ auto options = SelectKOptions(3, sort_keys);
+ auto expected = ChunkedArrayFromJSON(uint8(), {"[1, 1, 2]"});
+ this->Check(chunked_array, options, expected);
+ options.sort_keys[0].null_placement = NullPlacement::AtStart;
+ expected = ChunkedArrayFromJSON(uint8(), {"[null, null, 1]"});
+ this->Check(chunked_array, options, expected);
+}
+
+TYPED_TEST(TestSelectKWithChunkedArray, FullSelectKNull) {
+ auto chunked_array = ChunkedArrayFromJSON(uint8(), {
+ "[null, 1]",
+ "[3, null, 2]",
+ "[1]",
+ });
+ std::vector<SortKey> sort_keys{SortKey("a", SortOrder::Ascending)};
+ auto options = SelectKOptions(10, sort_keys);
+ options.sort_keys[0].null_placement = NullPlacement::AtStart;
+ auto expected = ChunkedArrayFromJSON(uint8(), {"[null, null, 1, 1, 2, 3]"});
+ this->Check(chunked_array, options, expected);
+ options.sort_keys[0].null_placement = NullPlacement::AtEnd;
+ expected = ChunkedArrayFromJSON(uint8(), {"[1, 1, 2, 3, null, null]"});
+ this->Check(chunked_array, options, expected);
+}
+
+TYPED_TEST(TestSelectKWithChunkedArray, PartialSelectKNullNaN) {
+ auto chunked_array = ChunkedArrayFromJSON(
+ float64(), {"[null, 1]", "[3, null, NaN]", "[10, NaN, 2]", "[1]"});
+ std::vector<SortKey> sort_keys{SortKey("a", SortOrder::Descending)};
+ auto options = SelectKOptions(3, sort_keys);
+ options.sort_keys[0].null_placement = NullPlacement::AtStart;
+ auto expected = ChunkedArrayFromJSON(float64(), {"[null, null, NaN]"});
+ this->Check(chunked_array, options, expected);
+ options.sort_keys[0].null_placement = NullPlacement::AtEnd;
+ expected = ChunkedArrayFromJSON(float64(), {"[10, 3, 2]"});
+ this->Check(chunked_array, options, expected);
+}
+
+TYPED_TEST(TestSelectKWithChunkedArray, FullSelectKNullNaN) {
+ auto chunked_array = ChunkedArrayFromJSON(
+ float64(), {"[null, 1]", "[3, null, NaN]", "[10, NaN, 2]", "[1]"});
+ std::vector<SortKey> sort_keys{SortKey("a", SortOrder::Descending)};
+ auto options = SelectKOptions(10, sort_keys);
+ options.sort_keys[0].null_placement = NullPlacement::AtStart;
+ // These check that nulls and Nan are sorted in a stable way, but do we want
that?
Review Comment:
I removed CheckIndices entirely
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
