Skip to content

GH-44729: [C++][Acero] Enable nested types for non-keys fields in AsofJoin operation #44871

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mroz45 wants to merge 5 commits into
base: main
Choose a base branch
from
+137 −52
Open

GH-44729: [C++][Acero] Enable nested types for non-keys fields in AsofJoin operation #44871

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3417cc3
update unmaterialized table
mroz45 Sep 19, 2024
ed354c0
a
mroz45 Oct 31, 2024
f3acf1a
Support for nested types non-key fields
mroz45 Nov 27, 2024
e8f7ccf
Test to Map data type
mroz45 Nov 28, 2024
9147b81
Merge branch 'main' into types
mroz45 Aug 27, 2026
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
4 changes: 4 additions & 0 deletions cpp/src/arrow/acero/asof_join_node.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(), " : ",
Expand Down
134 changes: 126 additions & 8 deletions cpp/src/arrow/acero/asof_join_node_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())}),
Expand Down Expand Up @@ -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
51 changes: 7 additions & 44 deletions cpp/src/arrow/acero/unmaterialized_table_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include <vector>
#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"
Expand Down Expand Up @@ -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 ",
Expand Down Expand Up @@ -165,45 +170,6 @@ class UnmaterializedCompositeTable {
num_rows += slice.Size();
}

template <class Type, class Builder = typename TypeTraits<Type>::BuilderType>
enable_if_boolean<Type, Status> static BuilderAppend(
Builder& builder, const std::shared_ptr<ArrayData>& source, uint64_t row) {
if (source->IsNull(row)) {
builder.UnsafeAppendNull();
return Status::OK();
}
builder.UnsafeAppend(bit_util::GetBit(source->template GetValues<uint8_t>(1), row));
return Status::OK();
}

template <class Type, class Builder = typename TypeTraits<Type>::BuilderType>
enable_if_t<is_fixed_width_type<Type>::value && !is_boolean_type<Type>::value,
Status> static BuilderAppend(Builder& builder,
const std::shared_ptr<ArrayData>& source,
uint64_t row) {
if (source->IsNull(row)) {
builder.UnsafeAppendNull();
return Status::OK();
}
using CType = typename TypeTraits<Type>::CType;
builder.UnsafeAppend(source->template GetValues<CType>(1)[row]);
return Status::OK();
}

template <class Type, class Builder = typename TypeTraits<Type>::BuilderType>
enable_if_base_binary<Type, Status> static BuilderAppend(
Builder& builder, const std::shared_ptr<ArrayData>& 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<offset_type>(1);
const offset_type offset0 = offsets[row];
const offset_type offset1 = offsets[row + 1];
return builder.Append(data + offset0, offset1 - offset0);
}

template <class Type, class Builder = typename arrow::TypeTraits<Type>::BuilderType>
arrow::Result<std::shared_ptr<arrow::Array>> materializeColumn(
const std::shared_ptr<arrow::DataType>& type, int i_col) {
Expand All @@ -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<Type, Builder>(
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());
Expand Down
Loading