diff --git a/src/iceberg/row/arrow_array_wrapper.cc b/src/iceberg/row/arrow_array_wrapper.cc index 0d6d48c15..5a9cb0894 100644 --- a/src/iceberg/row/arrow_array_wrapper.cc +++ b/src/iceberg/row/arrow_array_wrapper.cc @@ -134,7 +134,8 @@ class ArrowArrayStructLike::Impl { const ArrowArray* child_array = array_.get().children[pos]; const ArrowArrayView* child_view = array_view_.children[pos]; - return ExtractValue(child_schema, child_array, child_view, row_index_); + return ExtractValue(child_schema, child_array, child_view, + array_view_.offset + row_index_); } size_t num_fields() const { return static_cast(schema_.n_children); } @@ -252,8 +253,9 @@ class ArrowArrayArrayLike::Impl { array_.get().length); } - offset_ = ArrowArrayViewListChildOffset(&array_view_, row_index_); - length_ = ArrowArrayViewListChildOffset(&array_view_, row_index_ + 1) - offset_; + const int64_t row = array_view_.offset + row_index_; + offset_ = ArrowArrayViewListChildOffset(&array_view_, row); + length_ = ArrowArrayViewListChildOffset(&array_view_, row + 1) - offset_; return {}; } @@ -312,7 +314,7 @@ class ArrowArrayMapLike::Impl { const ArrowArrayView* keys_view = array_view_.children[0]->children[0]; return ExtractValue(keys_schema, keys_array, keys_view, - offset_ + static_cast(pos)); + EntryIndex(static_cast(pos))); } Result GetValue(size_t pos) const { @@ -326,7 +328,7 @@ class ArrowArrayMapLike::Impl { const ArrowArrayView* values_view = array_view_.children[0]->children[1]; return ExtractValue(values_schema, values_array, values_view, - offset_ + static_cast(pos)); + EntryIndex(static_cast(pos))); } size_t size() const { return static_cast(length_); } @@ -365,12 +367,17 @@ class ArrowArrayMapLike::Impl { // XXX: ArrowArrayViewListChildOffset does not work for map types. // We need to directly access the offsets buffer instead. auto* offsets_buffer = array_view_.buffer_views[1].data.as_int32; - offset_ = offsets_buffer[row_index_]; - length_ = offsets_buffer[row_index_ + 1] - offset_; + const int64_t row = array_view_.offset + row_index_; + offset_ = offsets_buffer[row]; + length_ = offsets_buffer[row + 1] - offset_; return {}; } + int64_t EntryIndex(int64_t pos) const { + return array_view_.children[0]->offset + offset_ + pos; + } + ArrowArrayView array_view_; internal::ArrowArrayViewGuard array_view_guard_{&array_view_}; diff --git a/src/iceberg/test/struct_like_test.cc b/src/iceberg/test/struct_like_test.cc index a57860116..e04d2c089 100644 --- a/src/iceberg/test/struct_like_test.cc +++ b/src/iceberg/test/struct_like_test.cc @@ -22,9 +22,12 @@ #include #include #include +#include #include #include +#include +#include #include #include #include @@ -663,4 +666,120 @@ TEST(ArrowArrayStructLike, Accessor) { } } +TEST(ArrowArrayStructLike, SlicedStruct) { + auto person_type = + ::arrow::struct_({::arrow::field("name", ::arrow::utf8(), /*nullable=*/false)}); + auto root_type = ::arrow::struct_( + {::arrow::field("id", ::arrow::int64(), /*nullable=*/false), + ::arrow::field("person", person_type, /*nullable=*/false), + ::arrow::field("tags", ::arrow::list(::arrow::int32()), /*nullable=*/false)}); + + auto arrow_array = ::arrow::json::ArrayFromJSONString(root_type, R"([ + {"id": 1, "person": {"name": "Alice"}, "tags": [1]}, + {"id": 2, "person": {"name": "Bob"}, "tags": [2, 3]}, + {"id": 3, "person": {"name": "Carol"}, "tags": []}])") + .ValueOrDie(); + auto sliced = arrow_array->Slice(1); + + ArrowSchema c_schema; + ArrowArray c_array; + internal::ArrowSchemaGuard schema_guard(&c_schema); + internal::ArrowArrayGuard array_guard(&c_array); + ASSERT_TRUE(::arrow::ExportType(*root_type, &c_schema).ok()); + ASSERT_TRUE(::arrow::ExportArray(*sliced, &c_array).ok()); + ASSERT_EQ(c_array.offset, 1); + + ICEBERG_UNWRAP_OR_FAIL(auto struct_like, ArrowArrayStructLike::Make(c_schema, c_array)); + + constexpr int64_t kNumRows = 2; + std::array ids = {2, 3}; + std::array names = {"Bob", "Carol"}; + std::array, kNumRows> tags = {std::vector{2, 3}, + std::vector{}}; + + for (int64_t i = 0; i < kNumRows; ++i) { + ASSERT_THAT(struct_like->Reset(i), IsOk()); + EXPECT_SCALAR_EQ(struct_like->GetField(0), int64_t, ids[i]); + + ICEBERG_UNWRAP_OR_FAIL(auto person_scalar, struct_like->GetField(1)); + auto person = std::get>(person_scalar); + EXPECT_SCALAR_EQ(person->GetField(0), std::string_view, names[i]); + + ICEBERG_UNWRAP_OR_FAIL(auto tags_scalar, struct_like->GetField(2)); + auto tags_like = std::get>(tags_scalar); + ASSERT_EQ(tags_like->size(), tags[i].size()); + for (size_t j = 0; j < tags[i].size(); ++j) { + EXPECT_SCALAR_EQ(tags_like->GetElement(j), int32_t, tags[i][j]); + } + } +} + +TEST(ArrowArrayStructLike, SlicedList) { + auto list_type = ::arrow::list(::arrow::int32()); + auto arrow_array = + ::arrow::json::ArrayFromJSONString(list_type, R"([[1], [2, 3], [4]])").ValueOrDie(); + auto sliced = arrow_array->Slice(1); + + ArrowSchema c_schema; + ArrowArray c_array; + internal::ArrowSchemaGuard schema_guard(&c_schema); + internal::ArrowArrayGuard array_guard(&c_array); + ASSERT_TRUE(::arrow::ExportType(*list_type, &c_schema).ok()); + ASSERT_TRUE(::arrow::ExportArray(*sliced, &c_array).ok()); + + ICEBERG_UNWRAP_OR_FAIL(auto array_like, ArrowArrayArrayLike::Make(c_schema, c_array)); + ASSERT_EQ(array_like->size(), 2); + EXPECT_SCALAR_EQ(array_like->GetElement(0), int32_t, 2); + EXPECT_SCALAR_EQ(array_like->GetElement(1), int32_t, 3); + + ASSERT_THAT(array_like->Reset(1), IsOk()); + ASSERT_EQ(array_like->size(), 1); + EXPECT_SCALAR_EQ(array_like->GetElement(0), int32_t, 4); +} + +TEST(ArrowArrayStructLike, SlicedMap) { + auto map_type = std::make_shared<::arrow::MapType>( + ::arrow::field("key", ::arrow::utf8(), /*nullable=*/false), + ::arrow::field("value", ::arrow::int32(), /*nullable=*/false)); + auto entries_source = + ::arrow::json::ArrayFromJSONString(map_type, R"([[["X", 0], ["Foo", 1], + ["Bar", 2], ["Baz", 3]]])") + .ValueOrDie(); + auto entries = + std::static_pointer_cast<::arrow::MapArray>(entries_source)->values()->Slice(1); + auto offsets = ::arrow::Buffer::FromVector(std::vector{0, 0, 2, 3, 3}); + auto map_data = ::arrow::ArrayData::Make(map_type, /*length=*/4, {nullptr, offsets}, + {entries->data()}, /*null_count=*/0); + auto sliced = ::arrow::MakeArray(map_data)->Slice(1); + ASSERT_TRUE(sliced->ValidateFull().ok()); + + ArrowSchema c_schema; + ArrowArray c_array; + internal::ArrowSchemaGuard schema_guard(&c_schema); + internal::ArrowArrayGuard array_guard(&c_array); + ASSERT_TRUE(::arrow::ExportType(*map_type, &c_schema).ok()); + ASSERT_TRUE(::arrow::ExportArray(*sliced, &c_array).ok()); + ASSERT_EQ(c_array.offset, 1); + ASSERT_EQ(c_array.children[0]->offset, 1); + + ICEBERG_UNWRAP_OR_FAIL(auto map_like, ArrowArrayMapLike::Make(c_schema, c_array)); + + constexpr int64_t kNumRows = 3; + std::array>, kNumRows> expected_maps = { + std::vector>{{"Foo", 1}, {"Bar", 2}}, + std::vector>{{"Baz", 3}}, + std::vector>{}, + }; + + for (int64_t i = 0; i < kNumRows; ++i) { + ASSERT_THAT(map_like->Reset(i), IsOk()); + const auto& expected_map = expected_maps[i]; + ASSERT_EQ(map_like->size(), expected_map.size()); + for (size_t j = 0; j < expected_map.size(); ++j) { + EXPECT_SCALAR_EQ(map_like->GetKey(j), std::string_view, expected_map[j].first); + EXPECT_SCALAR_EQ(map_like->GetValue(j), int32_t, expected_map[j].second); + } + } +} + } // namespace iceberg