diff --git a/cpp/src/arrow/compute/kernels/select_k_test.cc b/cpp/src/arrow/compute/kernels/select_k_test.cc index 47e4af58001c..3ed21dd4881c 100644 --- a/cpp/src/arrow/compute/kernels/select_k_test.cc +++ b/cpp/src/arrow/compute/kernels/select_k_test.cc @@ -415,6 +415,21 @@ TYPED_TEST(TestSelectKWithChunkedArrayTyped, RandomValuesWithSlices) { } } +TEST_F(TestSelectKWithChunkedArray, EmptyChunkedArray) { + auto chunked_array = std::make_shared(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{ "[null, 1]", @@ -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 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())}, diff --git a/cpp/src/arrow/compute/kernels/vector_select_k.cc b/cpp/src/arrow/compute/kernels/vector_select_k.cc index 7d94fa2ba330..7a63c76f4ef3 100644 --- a/cpp/src/arrow/compute/kernels/vector_select_k.cc +++ b/cpp/src/arrow/compute/kernels/vector_select_k.cc @@ -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()) { @@ -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()) {