Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions cpp/src/arrow/compute/kernels/select_k_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,21 @@ TYPED_TEST(TestSelectKWithChunkedArrayTyped, RandomValuesWithSlices) {
}
}

TEST_F(TestSelectKWithChunkedArray, EmptyChunkedArray) {
auto chunked_array = std::make_shared<ChunkedArray>(ArrayVector{}, uint8());
ASSERT_EQ(chunked_array->num_chunks(), 0);
ASSERT_EQ(chunked_array->length(), 0);

for (const auto& options :
{SelectKOptions::TopKDefault(3), SelectKOptions::BottomKDefault(3),
SelectKOptions::TopKDefault(0)}) {
ASSERT_OK_AND_ASSIGN(auto indices, SelectKUnstable(Datum(*chunked_array), options));
ASSERT_NE(indices, nullptr);
ValidateOutput(*indices);
ASSERT_EQ(indices->length(), 0);
}
}

TEST_F(TestSelectKWithChunkedArray, PartialSelectKNull) {
auto chunked_array = std::vector<std::string>{
"[null, 1]",
Expand Down Expand Up @@ -1103,6 +1118,25 @@ struct TestSelectKWithTable : public ::testing::Test {
}
};

TEST_F(TestSelectKWithTable, EmptyTable) {
auto schema = ::arrow::schema({
{field("a", uint8())},
{field("b", uint32())},
});
std::vector<std::string> input = {R"([])"};
auto table = TableFromJSON(schema, input);
ASSERT_EQ(table->num_rows(), 0);

for (const auto& options :
{SelectKOptions::TopKDefault(3, {"a"}), SelectKOptions::BottomKDefault(3, {"a"}),
SelectKOptions::TopKDefault(0, {"a"})}) {
ASSERT_OK_AND_ASSIGN(auto indices, SelectKUnstable(Datum(*table), options));
ASSERT_NE(indices, nullptr);
ValidateOutput(*indices);
ASSERT_EQ(indices->length(), 0);
}
}

TEST_F(TestSelectKWithTable, TopKOneColumnKey) {
auto schema = ::arrow::schema({
{field("a", uint8())},
Expand Down
6 changes: 6 additions & 0 deletions cpp/src/arrow/compute/kernels/vector_select_k.cc
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,9 @@ class ChunkedArraySelector : public TypeVisitor {

const auto num_chunks = chunked_array_.num_chunks();
if (num_chunks == 0) {
ARROW_ASSIGN_OR_RAISE(auto take_indices,
MakeMutableUInt64Array(0, ctx_->memory_pool()));
*output_ = Datum(take_indices);
return Status::OK();
}
if (k_ > chunked_array_.length()) {
Expand Down Expand Up @@ -624,6 +627,9 @@ class TableSelector : public TypeVisitor {

const auto num_rows = table_.num_rows();
if (num_rows == 0) {
ARROW_ASSIGN_OR_RAISE(auto take_indices,
MakeMutableUInt64Array(0, ctx_->memory_pool()));
*output_ = Datum(take_indices);
return Status::OK();
}
if (k_ > table_.num_rows()) {
Expand Down
Loading