From 3417cc3430d4e782d3f536e0d21a2d30cbf441bc Mon Sep 17 00:00:00 2001 From: kamilt Date: Thu, 19 Sep 2024 14:50:29 +0000 Subject: [PATCH 1/4] update unmaterialized table --- .../acero/unmaterialized_table_internal.h | 121 +++++++++++------- 1 file changed, 78 insertions(+), 43 deletions(-) diff --git a/cpp/src/arrow/acero/unmaterialized_table_internal.h b/cpp/src/arrow/acero/unmaterialized_table_internal.h index 86b1a763a603..fe6e7df899d5 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,7 @@ class UnmaterializedCompositeTable { MATERIALIZE_CASE(LARGE_STRING) MATERIALIZE_CASE(BINARY) MATERIALIZE_CASE(LARGE_BINARY) + MATERIALIZE_CASE(FIXED_SIZE_LIST) default: return arrow::Status::Invalid("Unsupported data type ", field->type()->ToString(), " for field ", @@ -165,44 +167,74 @@ 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> + 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> + enable_if_t::value, + Status> static BuilderAppend(Builder& builder, + const std::shared_ptr& source, + uint64_t row) { + //if (source->IsNull(row)) { + // builder.UnsafeAppendNull(); + // return Status::OK(); + //} + + return builder.AppendArraySlice(*source,row,1); + + //builder.Append(); + + //const int32_t list_size = internal::checked_cast(builder.type().get())->list_size(); + //builder.value_builder() + //FixedSizeListBuilder* + value_builder=internal::checked_cast(); + + //for(uint64_t i=0;i::CType; + //builder.UnsafeAppend(source->template GetValues(1)[row]); + //return Status::OK(); + }*/ template ::BuilderType> arrow::Result> materializeColumn( @@ -216,11 +248,14 @@ 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::Status st = builder.AppendArraySlice(*batch->column_data(column_index), + start, end - start); + ARROW_RETURN_NOT_OK(st); + // for (uint64_t rowNum = start; rowNum < end; ++rowNum) { + // arrow::Status st = BuilderAppend( + // builder, batch->column_data(column_index), rowNum); + // ARROW_RETURN_NOT_OK(st); + // } } else { for (uint64_t rowNum = start; rowNum < end; ++rowNum) { ARROW_RETURN_NOT_OK(builder.AppendNull()); From ed354c031a934da6b19f72142335b124e634eb75 Mon Sep 17 00:00:00 2001 From: kamilt Date: Thu, 31 Oct 2024 08:36:33 +0000 Subject: [PATCH 2/4] a --- .../acero/unmaterialized_table_internal.h | 74 +------------------ 1 file changed, 3 insertions(+), 71 deletions(-) diff --git a/cpp/src/arrow/acero/unmaterialized_table_internal.h b/cpp/src/arrow/acero/unmaterialized_table_internal.h index fe6e7df899d5..c0825efc8012 100644 --- a/cpp/src/arrow/acero/unmaterialized_table_internal.h +++ b/cpp/src/arrow/acero/unmaterialized_table_internal.h @@ -167,74 +167,7 @@ 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> - enable_if_t::value, - Status> static BuilderAppend(Builder& builder, - const std::shared_ptr& source, - uint64_t row) { - //if (source->IsNull(row)) { - // builder.UnsafeAppendNull(); - // return Status::OK(); - //} - - return builder.AppendArraySlice(*source,row,1); - - //builder.Append(); - - //const int32_t list_size = internal::checked_cast(builder.type().get())->list_size(); - //builder.value_builder() - //FixedSizeListBuilder* - value_builder=internal::checked_cast(); - - //for(uint64_t i=0;i::CType; - //builder.UnsafeAppend(source->template GetValues(1)[row]); - //return Status::OK(); - }*/ + template ::BuilderType> arrow::Result> materializeColumn( @@ -248,9 +181,8 @@ class UnmaterializedCompositeTable { for (const auto& unmaterialized_slice : slices) { const auto& [batch, start, end] = unmaterialized_slice.components[table_index]; if (batch) { - arrow::Status st = builder.AppendArraySlice(*batch->column_data(column_index), - start, end - start); - ARROW_RETURN_NOT_OK(st); +ARROW_RETURN_NOT_OK(builder.AppendArraySlice(*batch->column_data(column_index),start,end-start)); + // for (uint64_t rowNum = start; rowNum < end; ++rowNum) { // arrow::Status st = BuilderAppend( // builder, batch->column_data(column_index), rowNum); From f3acf1a97f6f190786511a98a7ce6324c1d5f3e0 Mon Sep 17 00:00:00 2001 From: kamilt Date: Wed, 27 Nov 2024 15:20:29 +0000 Subject: [PATCH 3/4] Support for nested types non-key fields --- cpp/src/arrow/acero/asof_join_node.cc | 3 + cpp/src/arrow/acero/asof_join_node_test.cc | 104 ++++++++++++++++-- .../acero/unmaterialized_table_internal.h | 13 +-- 3 files changed, 103 insertions(+), 17 deletions(-) diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index 55fa45543e4f..2f1947c94b2f 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -1229,6 +1229,9 @@ 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: 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 59a9b4ebba12..d0b6e75ff210 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,101 @@ 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); +} } // 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 c0825efc8012..cf382b468395 100644 --- a/cpp/src/arrow/acero/unmaterialized_table_internal.h +++ b/cpp/src/arrow/acero/unmaterialized_table_internal.h @@ -114,6 +114,8 @@ class UnmaterializedCompositeTable { MATERIALIZE_CASE(BINARY) MATERIALIZE_CASE(LARGE_BINARY) MATERIALIZE_CASE(FIXED_SIZE_LIST) + MATERIALIZE_CASE(LIST) + MATERIALIZE_CASE(STRUCT) default: return arrow::Status::Invalid("Unsupported data type ", field->type()->ToString(), " for field ", @@ -167,8 +169,6 @@ class UnmaterializedCompositeTable { num_rows += slice.Size(); } - - template ::BuilderType> arrow::Result> materializeColumn( const std::shared_ptr& type, int i_col) { @@ -181,13 +181,8 @@ class UnmaterializedCompositeTable { for (const auto& unmaterialized_slice : slices) { const auto& [batch, start, end] = unmaterialized_slice.components[table_index]; if (batch) { -ARROW_RETURN_NOT_OK(builder.AppendArraySlice(*batch->column_data(column_index),start,end-start)); - - // 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()); From e8f7ccff2e8f97bb6fd21feaab84bd284a1686f1 Mon Sep 17 00:00:00 2001 From: kamilt Date: Thu, 28 Nov 2024 10:48:59 +0000 Subject: [PATCH 4/4] Test to Map data type --- cpp/src/arrow/acero/asof_join_node.cc | 1 + cpp/src/arrow/acero/asof_join_node_test.cc | 30 +++++++++++++++++++ .../acero/unmaterialized_table_internal.h | 1 + 3 files changed, 32 insertions(+) diff --git a/cpp/src/arrow/acero/asof_join_node.cc b/cpp/src/arrow/acero/asof_join_node.cc index 2f1947c94b2f..7b17594428b6 100644 --- a/cpp/src/arrow/acero/asof_join_node.cc +++ b/cpp/src/arrow/acero/asof_join_node.cc @@ -1232,6 +1232,7 @@ class AsofJoinNode : public ExecNode { 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 d0b6e75ff210..43c7c6845d0b 100644 --- a/cpp/src/arrow/acero/asof_join_node_test.cc +++ b/cpp/src/arrow/acero/asof_join_node_test.cc @@ -1912,5 +1912,35 @@ TEST(AsofJoinTest, StructTestDataType) { ])"); 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 cf382b468395..6b953f25cb0f 100644 --- a/cpp/src/arrow/acero/unmaterialized_table_internal.h +++ b/cpp/src/arrow/acero/unmaterialized_table_internal.h @@ -116,6 +116,7 @@ class UnmaterializedCompositeTable { 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 ",