diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index 3970050e502..89dfeba21f7 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -1231,6 +1231,10 @@ class AsofJoinNode : public ExecNode { case Type::LARGE_STRING: case Type::BINARY: case Type::LARGE_BINARY: + case Type::LIST: + case Type::FIXED_SIZE_LIST: + case Type::STRUCT: + case Type::MAP: return Status::OK(); default: return Status::Invalid("Unsupported type for data field ", field->name(), " : ", diff --git a/cpp/src/arrow/acero/asof_join_node_test.cc b/cpp/src/arrow/acero/asof_join_node_test.cc index 59a9b4ebba1..43c7c6845d0 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1294,14 +1294,6 @@ TRACED_TEST(AsofJoinTest, TestUnsupportedByType, { field("r0_v0", float32())})); }) -TRACED_TEST(AsofJoinTest, TestUnsupportedDatatype, { - // List is unsupported - DoRunInvalidTypeTest( - schema({field("time", int64()), field("key", int32()), field("l_v0", float64())}), - schema({field("time", int64()), field("key", int32()), - field("r0_v0", list(int32()))})); -}) - TRACED_TEST(AsofJoinTest, TestMissingKeys, { DoRunMissingKeysTest( schema({field("time1", int64()), field("key", int32()), field("l_v0", float64())}), @@ -1824,5 +1816,131 @@ TEST(AsofJoinTest, OneSideTsAllGreaterThanTheOther) { } } +// GH-44729: Testing nested data type for non-key fields +TEST(AsofJoinTest, FixedListDataType) { + const int32_t list_size = 3; + auto list_type = arrow::fixed_size_list(arrow::int32(), list_size); + + auto left_batch = ExecBatchFromJSON({int64()}, R"([[1], [2], [3]])"); + auto right_batch = ExecBatchFromJSON({list_type, int64()}, R"([ + [[0, 1, 2], 2], + [[3, 4, 5], 3], + [[6, 7, 8], 4] + ])"); + + Declaration left{"exec_batch_source", + ExecBatchSourceNodeOptions(schema({field("on", int64())}), + {std::move(left_batch)})}; + Declaration right{"exec_batch_source", + ExecBatchSourceNodeOptions( + schema({field("colVals", list_type), field("on", int64())}), + {std::move(right_batch)})}; + + AsofJoinNodeOptions asof_join_opts({{{"on"}, {}}, {{"on"}, {}}}, 1); + Declaration asof_join{ + "asofjoin", {std::move(left), std::move(right)}, std::move(asof_join_opts)}; + + ASSERT_OK_AND_ASSIGN(auto result, DeclarationToExecBatches(std::move(asof_join))); + + auto exp_batch = ExecBatchFromJSON({int64(), list_type}, R"([ + [1, [0, 1, 2]], + [2, [0, 1, 2]], + [3, [3, 4, 5]] + ])"); + + AssertExecBatchesEqual(result.schema, {exp_batch}, result.batches); +} + +TEST(AsofJoinTest, ListDataType) { + auto list_type = list(int32()); + + auto left_batch = ExecBatchFromJSON({int64()}, R"([[1], [2], [3]])"); + auto right_batch = ExecBatchFromJSON({list_type, int64()}, R"([ + [[0, 1, 2, 9], 2], + [[3, 4, 5, 7], 3], + [[6, 7, 8], 4] + ])"); + + Declaration left{"exec_batch_source", + ExecBatchSourceNodeOptions(schema({field("on", int64())}), + {std::move(left_batch)})}; + Declaration right{"exec_batch_source", + ExecBatchSourceNodeOptions( + schema({field("colVals", list_type), field("on", int64())}), + {std::move(right_batch)})}; + + AsofJoinNodeOptions asof_join_opts({{{"on"}, {}}, {{"on"}, {}}}, 1); + Declaration asof_join{ + "asofjoin", {std::move(left), std::move(right)}, std::move(asof_join_opts)}; + + ASSERT_OK_AND_ASSIGN(auto result, DeclarationToExecBatches(std::move(asof_join))); + auto exp_batch = ExecBatchFromJSON({int64(), list_type}, R"([ + [1, [0, 1, 2, 9]], + [2, [0, 1, 2, 9]], + [3, [3, 4, 5, 7]] + ])"); + + AssertExecBatchesEqual(result.schema, {exp_batch}, result.batches); +} + +TEST(AsofJoinTest, StructTestDataType) { + auto struct_type = struct_({field("key", utf8()), field("value", int64())}); + + auto left_batch = ExecBatchFromJSON({int64()}, R"([[1], [2], [3]])"); + auto right_batch = ExecBatchFromJSON({struct_type, int64()}, R"([ + [{"key": "a", "value": 1}, 2], + [{"key": "b", "value": 3}, 3], + [{"key": "c", "value": 5}, 4] + ])"); + + Declaration left{"exec_batch_source", + ExecBatchSourceNodeOptions(schema({field("on", int64())}), + {std::move(left_batch)})}; + Declaration right{"exec_batch_source", + ExecBatchSourceNodeOptions( + schema({field("col", struct_type), field("on", int64())}), + {std::move(right_batch)})}; + AsofJoinNodeOptions asof_join_opts({{{"on"}, {}}, {{"on"}, {}}}, 1); + Declaration asof_join{ + "asofjoin", {std::move(left), std::move(right)}, std::move(asof_join_opts)}; + ASSERT_OK_AND_ASSIGN(auto result, DeclarationToExecBatches(std::move(asof_join))); + + auto exp_batch = ExecBatchFromJSON({int64(), struct_type}, R"([ + [1, {"key": "a", "value": 1}], + [2, {"key": "a", "value": 1}], + [3, {"key": "b", "value": 3}] + ])"); + AssertExecBatchesEqual(result.schema, {exp_batch}, result.batches); +} + +TEST(AsofJoinTest, MapTestDataType) { + auto map_type = map(int64(), int64()); + + auto left_batch = ExecBatchFromJSON({int64()}, R"([[1], [2], [3]])"); + auto right_batch = ExecBatchFromJSON({map_type, int64()}, R"([ + [[[11, 111], [22, 222]], 2], + [[[33, 333], [44, 444], [77, 777]], 3], + [[[55, 555], [66, 666]], 4] + ])"); + + Declaration left{"exec_batch_source", + ExecBatchSourceNodeOptions(schema({field("on", int64())}), + {std::move(left_batch)})}; + Declaration right{ + "exec_batch_source", + ExecBatchSourceNodeOptions(schema({field("col", map_type), field("on", int64())}), + {std::move(right_batch)})}; + AsofJoinNodeOptions asof_join_opts({{{"on"}, {}}, {{"on"}, {}}}, 1); + Declaration asof_join{ + "asofjoin", {std::move(left), std::move(right)}, std::move(asof_join_opts)}; + + ASSERT_OK_AND_ASSIGN(auto result, DeclarationToExecBatches(std::move(asof_join))); + auto exp_batch = ExecBatchFromJSON({int64(), map_type}, R"([ + [1, [[11, 111], [22, 222]]], + [2, [[11, 111], [22, 222]]], + [3, [[33, 333], [44, 444], [77, 777]]] + ])"); + AssertExecBatchesEqual(result.schema, {exp_batch}, result.batches); +} } // namespace acero } // namespace arrow diff --git a/cpp/src/arrow/acero/unmaterialized_table_internal.h b/cpp/src/arrow/acero/unmaterialized_table_internal.h index 86b1a763a60..6b953f25cb0 100644 --- a/cpp/src/arrow/acero/unmaterialized_table_internal.h +++ b/cpp/src/arrow/acero/unmaterialized_table_internal.h @@ -21,6 +21,7 @@ #include #include "arrow/array/builder_base.h" #include "arrow/array/builder_binary.h" +#include "arrow/array/builder_nested.h" #include "arrow/array/builder_primitive.h" #include "arrow/memory_pool.h" #include "arrow/record_batch.h" @@ -112,6 +113,10 @@ class UnmaterializedCompositeTable { MATERIALIZE_CASE(LARGE_STRING) MATERIALIZE_CASE(BINARY) MATERIALIZE_CASE(LARGE_BINARY) + MATERIALIZE_CASE(FIXED_SIZE_LIST) + MATERIALIZE_CASE(LIST) + MATERIALIZE_CASE(STRUCT) + MATERIALIZE_CASE(MAP) default: return arrow::Status::Invalid("Unsupported data type ", field->type()->ToString(), " for field ", @@ -165,45 +170,6 @@ class UnmaterializedCompositeTable { num_rows += slice.Size(); } - template ::BuilderType> - enable_if_boolean static BuilderAppend( - Builder& builder, const std::shared_ptr& source, uint64_t row) { - if (source->IsNull(row)) { - builder.UnsafeAppendNull(); - return Status::OK(); - } - builder.UnsafeAppend(bit_util::GetBit(source->template GetValues(1), row)); - return Status::OK(); - } - - template ::BuilderType> - enable_if_t::value && !is_boolean_type::value, - Status> static BuilderAppend(Builder& builder, - const std::shared_ptr& source, - uint64_t row) { - if (source->IsNull(row)) { - builder.UnsafeAppendNull(); - return Status::OK(); - } - using CType = typename TypeTraits::CType; - builder.UnsafeAppend(source->template GetValues(1)[row]); - return Status::OK(); - } - - template ::BuilderType> - enable_if_base_binary static BuilderAppend( - Builder& builder, const std::shared_ptr& source, uint64_t row) { - if (source->IsNull(row)) { - return builder.AppendNull(); - } - using offset_type = typename Type::offset_type; - const uint8_t* data = source->buffers[2]->data(); - const offset_type* offsets = source->GetValues(1); - const offset_type offset0 = offsets[row]; - const offset_type offset1 = offsets[row + 1]; - return builder.Append(data + offset0, offset1 - offset0); - } - template ::BuilderType> arrow::Result> materializeColumn( const std::shared_ptr& type, int i_col) { @@ -216,11 +182,8 @@ class UnmaterializedCompositeTable { for (const auto& unmaterialized_slice : slices) { const auto& [batch, start, end] = unmaterialized_slice.components[table_index]; if (batch) { - for (uint64_t rowNum = start; rowNum < end; ++rowNum) { - arrow::Status st = BuilderAppend( - builder, batch->column_data(column_index), rowNum); - ARROW_RETURN_NOT_OK(st); - } + ARROW_RETURN_NOT_OK(builder.AppendArraySlice(*batch->column_data(column_index), + start, end - start)); } else { for (uint64_t rowNum = start; rowNum < end; ++rowNum) { ARROW_RETURN_NOT_OK(builder.AppendNull());