From cc177062df2ef20b433733b1232cfd5ff97d157e Mon Sep 17 00:00:00 2001 From: Oleksii Gubanov <153339237+0guban0v@users.noreply.github.com> Date: Thu, 6 Aug 2026 21:53:33 -0700 Subject: [PATCH 1/3] Add Arrow fixed-size-list ingress --- cpp/src/interop/arrow_data_structures.cpp | 23 +++ cpp/src/interop/arrow_utilities.cpp | 22 ++- cpp/src/interop/arrow_utilities.hpp | 21 +- cpp/src/interop/from_arrow_device.cu | 41 +++- cpp/src/interop/from_arrow_host.cu | 58 +++++- cpp/src/interop/from_arrow_host.hpp | 19 +- .../interop/arrow_data_structures_test.cpp | 49 ++++- cpp/tests/interop/from_arrow_device_test.cpp | 119 ++++++++++++ cpp/tests/interop/from_arrow_host_test.cpp | 181 ++++++++++++++++++ cpp/tests/interop/from_arrow_stream_test.cpp | 91 ++++++++- python/pylibcudf/tests/test_table.py | 13 ++ 11 files changed, 622 insertions(+), 15 deletions(-) diff --git a/cpp/src/interop/arrow_data_structures.cpp b/cpp/src/interop/arrow_data_structures.cpp index 16a32d638a06..41f0b1ac4030 100644 --- a/cpp/src/interop/arrow_data_structures.cpp +++ b/cpp/src/interop/arrow_data_structures.cpp @@ -21,6 +21,23 @@ namespace cudf::interop { +namespace { + +bool contains_fixed_size_list(ArrowSchema const& schema) +{ + ArrowSchemaView schema_view; + NANOARROW_THROW_NOT_OK(ArrowSchemaViewInit(&schema_view, &schema, nullptr)); + if (schema_view.type == NANOARROW_TYPE_FIXED_SIZE_LIST) { return true; } + + for (auto i = 0; i < schema.n_children; ++i) { + if (contains_fixed_size_list(*schema.children[i])) { return true; } + } + + return schema.dictionary != nullptr && contains_fixed_size_list(*schema.dictionary); +} + +} // namespace + /** * @brief A wrapper around ArrowDeviceArray data used for flexible lifetime management. * @@ -66,6 +83,12 @@ struct arrow_array_container { rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { + // TODO: Remove this guard when owning wrappers can export fixed-size-list with matching + // schema and buffers for every target device. + CUDF_EXPECTS(!contains_fixed_size_list(schema_), + "Owning Arrow device wrappers do not support fixed-size-list input until " + "fixed-size-list egress is implemented", + cudf::data_type_error); switch (input_.device_type) { case ARROW_DEVICE_CUDA: case ARROW_DEVICE_CUDA_HOST: diff --git a/cpp/src/interop/arrow_utilities.cpp b/cpp/src/interop/arrow_utilities.cpp index eb4a0bb8851a..ec7bab8d6aae 100644 --- a/cpp/src/interop/arrow_utilities.cpp +++ b/cpp/src/interop/arrow_utilities.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,6 +11,8 @@ #include +#include + namespace cudf { namespace detail { data_type arrow_to_cudf_type(ArrowSchemaView const* arrow_view) @@ -33,7 +35,8 @@ data_type arrow_to_cudf_type(ArrowSchemaView const* arrow_view) case NANOARROW_TYPE_STRING_VIEW: case NANOARROW_TYPE_LARGE_STRING: return data_type(type_id::STRING); case NANOARROW_TYPE_LIST: - case NANOARROW_TYPE_LARGE_LIST: return data_type(type_id::LIST); + case NANOARROW_TYPE_LARGE_LIST: + case NANOARROW_TYPE_FIXED_SIZE_LIST: return data_type(type_id::LIST); case NANOARROW_TYPE_DICTIONARY: return data_type(type_id::DICTIONARY32); case NANOARROW_TYPE_STRUCT: return data_type(type_id::STRUCT); case NANOARROW_TYPE_TIMESTAMP: { @@ -62,6 +65,21 @@ data_type arrow_to_cudf_type(ArrowSchemaView const* arrow_view) } } +bool is_fixed_size_list(ArrowSchemaView const* arrow_view) +{ + return arrow_view->type == NANOARROW_TYPE_FIXED_SIZE_LIST; +} + +size_type fixed_size_list_width(ArrowSchemaView const* arrow_view) +{ + CUDF_EXPECTS( + is_fixed_size_list(arrow_view), "Expected a fixed-size-list schema", cudf::data_type_error); + CUDF_EXPECTS(arrow_view->fixed_size >= 0, + "fixed-size-list width must be non-negative", + std::invalid_argument); + return static_cast(arrow_view->fixed_size); +} + ArrowType id_to_arrow_type(cudf::type_id id) { switch (id) { diff --git a/cpp/src/interop/arrow_utilities.hpp b/cpp/src/interop/arrow_utilities.hpp index 09e091ddc130..93c65761796b 100644 --- a/cpp/src/interop/arrow_utilities.hpp +++ b/cpp/src/interop/arrow_utilities.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -28,6 +28,25 @@ static constexpr int fixed_width_data_buffer_idx = 1; */ data_type arrow_to_cudf_type(ArrowSchemaView const* arrow_view); +/** + * @brief Check whether the given schema view describes an Arrow fixed-size-list + * + * @param arrow_view SchemaView to check + * @return True if the schema describes a fixed-size-list + */ +bool is_fixed_size_list(ArrowSchemaView const* arrow_view); + +/** + * @brief Return the number of child elements per row of a fixed-size-list schema + * + * @throw cudf::data_type_error if `arrow_view` is not a fixed-size-list + * @throw std::invalid_argument if the declared width is negative + * + * @param arrow_view SchemaView to pull the fixed size from + * @return Number of child elements per row + */ +size_type fixed_size_list_width(ArrowSchemaView const* arrow_view); + /** * @brief Map cudf column type id to ArrowType id * diff --git a/cpp/src/interop/from_arrow_device.cu b/cpp/src/interop/from_arrow_device.cu index aed25a2c5496..af65503b9b78 100644 --- a/cpp/src/interop/from_arrow_device.cu +++ b/cpp/src/interop/from_arrow_device.cu @@ -4,6 +4,7 @@ */ #include "arrow_utilities.hpp" +#include "from_arrow_host.hpp" #include #include @@ -324,12 +325,6 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()( size_type const num_rows = input->length; size_type const offset = input->offset; size_type const null_count = input->null_count; - auto offsets_view = column_view{data_type(type_id::INT32), - (num_rows == 0) ? 0 : (offset + num_rows + 1), - input->buffers[fixed_width_data_buffer_idx], - nullptr, - 0, - 0}; ArrowSchemaView child_schema_view; NANOARROW_THROW_NOT_OK( @@ -341,8 +336,38 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()( // in the scenario where we were sliced and there are more elements in the child_view // than can be referenced by the sliced offsets, we need to slice the child_view // so that when `get_sliced_child` is called, we still produce the right result - auto max_child_offset = - num_rows == 0 ? 0 : cudf::detail::get_value(offsets_view, offset + num_rows, stream); + column_view offsets_view; + size_type max_child_offset = 0; + if (is_fixed_size_list(schema)) { + // fixed-size-list arrays carry no offsets buffer, so synthesize {0, w, 2w, ...}. + // these are absolute rather than normalized because the outer column_view applies a + // single offset to both the null mask and the children. + auto const width = static_cast(fixed_size_list_width(schema)); + auto const num_offsets = static_cast(offset) + num_rows + 1; + auto const child_end = (num_offsets - 1) * width; + constexpr auto max_size = static_cast(std::numeric_limits::max()); + CUDF_EXPECTS(num_offsets <= max_size && child_end <= max_size, + "fixed-size-list offsets exceed cuDF's maximum supported row count " + "(cudf::size_type).", + std::overflow_error); + max_child_offset = (num_rows == 0) ? 0 : static_cast(child_end); + if (num_rows == 0) { + offsets_view = column_view{data_type{type_id::INT32}, 0, nullptr, nullptr, 0, 0}; + } else { + owned.emplace_back(make_fixed_size_list_offsets( + static_cast(num_offsets), static_cast(width), stream, mr)); + offsets_view = owned.back()->view(); + } + } else { + offsets_view = column_view{data_type(type_id::INT32), + (num_rows == 0) ? 0 : (offset + num_rows + 1), + input->buffers[fixed_width_data_buffer_idx], + nullptr, + 0, + 0}; + max_child_offset = + num_rows == 0 ? 0 : cudf::detail::get_value(offsets_view, offset + num_rows, stream); + } child_view = cudf::slice(child_view, {0, max_child_offset}, stream).front(); return std::make_tuple( diff --git a/cpp/src/interop/from_arrow_host.cu b/cpp/src/interop/from_arrow_host.cu index 8f6cdf4483e2..89b57ecb1e0a 100644 --- a/cpp/src/interop/from_arrow_host.cu +++ b/cpp/src/interop/from_arrow_host.cu @@ -30,6 +30,8 @@ #include #include +#include + #include #include #include @@ -284,6 +286,37 @@ std::unique_ptr dispatch_copy_from_arrow_host::operator()length, std::move(child_columns), null_count, std::move(*out_mask), stream, mr); } +/** + * @brief Synthesize the offsets column and child bounds for a fixed-size-list array + * + * Mirrors the (offsets, child-offset, child-length) contract of `get_offsets_column`. + * Fixed-size-list arrays carry no offsets buffer, so `buffers[fixed_width_data_buffer_idx]` + * is never read here. The returned offsets are normalized to start at zero, matching + * `copy_offsets_column`; the absolute start of the child range is returned separately. + */ +std::tuple, int64_t, int64_t> get_fixed_size_list_offsets( + ArrowSchemaView const* schema, + ArrowArray const* input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto const width = static_cast(fixed_size_list_width(schema)); + auto const offset = input->offset * width; + auto const length = input->length * width; + + constexpr auto max_size = static_cast(std::numeric_limits::max()); + CUDF_EXPECTS(length <= max_size, + "Number of fixed-size-list child elements exceeds cuDF's maximum supported " + "row count (cudf::size_type).", + std::overflow_error); + + return std::tuple{ + make_fixed_size_list_offsets( + static_cast(input->length) + 1, static_cast(width), stream, mr), + offset, + length}; +} + template <> std::unique_ptr dispatch_copy_from_arrow_host::operator()( ArrowSchemaView const* schema, ArrowArray const* input, data_type type, bool skip_mask) @@ -293,7 +326,9 @@ std::unique_ptr dispatch_copy_from_arrow_host::operator()schema->children[0], nullptr)); @@ -394,9 +429,30 @@ std::tuple, int64_t, int64_t> copy_offsets_column( } // namespace +std::unique_ptr make_fixed_size_list_offsets(size_type size, + size_type width, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto offsets = + make_numeric_column(data_type{type_id::INT32}, size, mask_state::UNALLOCATED, stream, mr); + auto d_offsets = offsets->mutable_view().begin(); + thrust::sequence(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + d_offsets, + d_offsets + size, + size_type{0}, + width); + return offsets; +} + /** * @brief Utility to copy the offsets from the given input (strings or list) to a * cudf column + * + * @note This requires `input` to carry an offsets buffer at + * `fixed_width_data_buffer_idx`, which it reads before inspecting `schema->type`. + * Fixed-size-list arrays have no offsets buffer (`n_buffers == 1`), so they must be + * routed to `get_fixed_size_list_offsets` instead. */ std::tuple, int64_t, int64_t> get_offsets_column( ArrowSchemaView const* schema, diff --git a/cpp/src/interop/from_arrow_host.hpp b/cpp/src/interop/from_arrow_host.hpp index 37d50a86084d..c571230bba69 100644 --- a/cpp/src/interop/from_arrow_host.hpp +++ b/cpp/src/interop/from_arrow_host.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -48,5 +48,22 @@ std::tuple, int64_t, int64_t> get_offsets_column( rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); +/** + * @brief Create the offsets column for a fixed-size-list array + * + * Arrow fixed-size-list arrays carry no offsets buffer; the offsets are implicit. + * This generates `size` offsets of the form `{0, width, 2*width, ...}`. + * + * @param size Number of offsets to generate (normally num_rows + 1) + * @param width Number of child elements per list row + * @param stream CUDA stream used for device memory operations + * @param mr Device memory resource to use for all device memory allocations + * @return INT32 offsets column + */ +std::unique_ptr make_fixed_size_list_offsets(size_type size, + size_type width, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + } // namespace detail } // namespace cudf diff --git a/cpp/tests/interop/arrow_data_structures_test.cpp b/cpp/tests/interop/arrow_data_structures_test.cpp index 36176164fdcb..1cc1ff44794e 100644 --- a/cpp/tests/interop/arrow_data_structures_test.cpp +++ b/cpp/tests/interop/arrow_data_structures_test.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -19,6 +19,33 @@ struct ArrowColumnTest : public cudf::test::BaseFixture {}; +namespace { + +nanoarrow::UniqueSchema make_fixed_size_list_schema(bool wrap_in_struct) +{ + nanoarrow::UniqueSchema schema; + ArrowSchemaInit(schema.get()); + + auto* list_schema = schema.get(); + if (wrap_in_struct) { + NANOARROW_THROW_NOT_OK(ArrowSchemaSetTypeStruct(schema.get(), 1)); + list_schema = schema->children[0]; + } + + NANOARROW_THROW_NOT_OK( + ArrowSchemaSetTypeFixedSize(list_schema, NANOARROW_TYPE_FIXED_SIZE_LIST, 3)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetType(list_schema->children[0], NANOARROW_TYPE_INT64)); + return schema; +} + +ArrowDeviceArray make_empty_device_array() +{ + return ArrowDeviceArray{ + .array = {}, .device_id = 0, .device_type = ARROW_DEVICE_CUDA, .sync_event = nullptr}; +} + +} // namespace + template auto export_to_arrow(T& obj, ArrowDeviceType device_type = ARROW_DEVICE_CUDA) { @@ -166,6 +193,16 @@ TEST_F(ArrowColumnTest, ToFromHost) CUDF_TEST_EXPECT_COLUMNS_EQUAL(int_col, arrow_column_from_arrow_array.view()); } +TEST_F(ArrowColumnTest, FixedSizeListDeviceInputRejected) +{ + auto schema = make_fixed_size_list_schema(false); + auto array = make_empty_device_array(); + + EXPECT_THROW( + { static_cast(cudf::interop::arrow_column(std::move(*schema.get()), std::move(array))); }, + cudf::data_type_error); +} + struct ArrowTableTest : public cudf::test::BaseFixture {}; TEST_F(ArrowTableTest, TwoWayConversion) @@ -284,3 +321,13 @@ TEST_F(ArrowTableTest, FromArrowArrayStream) auto result = cudf::interop::arrow_table(std::move(stream)); CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), result.view()); } + +TEST_F(ArrowTableTest, NestedFixedSizeListDeviceInputRejected) +{ + auto schema = make_fixed_size_list_schema(true); + auto array = make_empty_device_array(); + + EXPECT_THROW( + { static_cast(cudf::interop::arrow_table(std::move(*schema.get()), std::move(array))); }, + cudf::data_type_error); +} diff --git a/cpp/tests/interop/from_arrow_device_test.cpp b/cpp/tests/interop/from_arrow_device_test.cpp index 5de19ba0a9ac..eb8f899fa089 100644 --- a/cpp/tests/interop/from_arrow_device_test.cpp +++ b/cpp/tests/interop/from_arrow_device_test.cpp @@ -268,6 +268,125 @@ TEST_F(FromArrowDeviceTest, NestedList) CUDF_TEST_EXPECT_TABLES_EQUAL(*got_cudf_table_view, from_struct); } +namespace { + +// A fixed_size_list array has only a validity buffer; there is no offsets buffer to wire up. +void populate_fixed_size_list_from_col(ArrowArray* arr, cudf::lists_column_view view) +{ + arr->length = view.size(); + arr->null_count = view.null_count(); + + NANOARROW_THROW_NOT_OK(ArrowBufferSetAllocator(ArrowArrayBuffer(arr, 0), noop_alloc)); + ArrowArrayValidityBitmap(arr)->buffer.size_bytes = + cudf::bitmask_allocation_size_bytes(view.size()); + ArrowArrayValidityBitmap(arr)->buffer.data = + const_cast(reinterpret_cast(view.null_mask())); +} + +// ArrowSchemaInitFromType does not support NANOARROW_TYPE_FIXED_SIZE_LIST (no format +// template, returns EINVAL); ArrowSchemaSetTypeFixedSize is the supported path and leaves +// the allocated "item" child with a NULL format, so the child type is set explicitly. +nanoarrow::UniqueSchema make_fixed_size_list_device_schema(int32_t width) +{ + nanoarrow::UniqueSchema schema; + ArrowSchemaInit(schema.get()); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetTypeStruct(schema.get(), 1)); + + NANOARROW_THROW_NOT_OK( + ArrowSchemaSetTypeFixedSize(schema->children[0], NANOARROW_TYPE_FIXED_SIZE_LIST, width)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0], "a")); + schema->children[0]->flags = 0; + + NANOARROW_THROW_NOT_OK( + ArrowSchemaSetType(schema->children[0]->children[0], NANOARROW_TYPE_INT64)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0]->children[0], "element")); + schema->children[0]->children[0]->flags = 0; + + return schema; +} + +} // namespace + +TEST_F(FromArrowDeviceTest, FixedSizeListColumn) +{ + auto col = + cudf::test::lists_column_wrapper{{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; + cudf::table_view expected_table_view({col}); + + auto input_schema = make_fixed_size_list_device_schema(3); + + nanoarrow::UniqueArray input_array; + EXPECT_EQ(NANOARROW_OK, ArrowArrayInitFromSchema(input_array.get(), input_schema.get(), nullptr)); + input_array->length = expected_table_view.num_rows(); + auto top_list = input_array->children[0]; + cudf::lists_column_view lview{expected_table_view.column(0)}; + populate_fixed_size_list_from_col(top_list, lview); + populate_from_col(top_list->children[0], lview.child()); + NANOARROW_THROW_NOT_OK( + ArrowArrayFinishBuilding(input_array.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); + + ArrowDeviceArray input_device_array; + input_device_array.device_id = rmm::get_current_cuda_device().value(); + input_device_array.device_type = ARROW_DEVICE_CUDA; + input_device_array.sync_event = nullptr; + memcpy(&input_device_array.array, input_array.get(), sizeof(ArrowArray)); + + auto got_cudf_table_view = cudf::from_arrow_device(input_schema.get(), &input_device_array); + EXPECT_EQ(got_cudf_table_view->column(0).type(), cudf::data_type{cudf::type_id::LIST}); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected_table_view, *got_cudf_table_view); + + ArrowDeviceArray direct_device_array; + memcpy(&direct_device_array.array, input_array->children[0], sizeof(ArrowArray)); + direct_device_array.device_id = input_device_array.device_id; + direct_device_array.device_type = ARROW_DEVICE_CUDA; + direct_device_array.sync_event = nullptr; + auto got_direct_col = + cudf::from_arrow_device_column(input_schema->children[0], &direct_device_array); + EXPECT_EQ(got_direct_col->type(), cudf::data_type{cudf::type_id::LIST}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_table_view.column(0), *got_direct_col); + + auto got_cudf_col = cudf::from_arrow_device_column(input_schema.get(), &input_device_array); + EXPECT_EQ(got_cudf_col->type(), cudf::data_type{cudf::type_id::STRUCT}); + cudf::table_view from_struct{ + std::vector(got_cudf_col->child_begin(), got_cudf_col->child_end())}; + CUDF_TEST_EXPECT_TABLES_EQUAL(*got_cudf_table_view, from_struct); +} + +TEST_F(FromArrowDeviceTest, FixedSizeListColumnSliced) +{ + auto col = + cudf::test::lists_column_wrapper{{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; + cudf::table_view full_table_view({col}); + auto expected = cudf::slice(full_table_view.column(0), {1, 3}).front(); + cudf::table_view expected_table_view({expected}); + + auto input_schema = make_fixed_size_list_device_schema(3); + + nanoarrow::UniqueArray input_array; + EXPECT_EQ(NANOARROW_OK, ArrowArrayInitFromSchema(input_array.get(), input_schema.get(), nullptr)); + auto top_list = input_array->children[0]; + cudf::lists_column_view lview{full_table_view.column(0)}; + populate_fixed_size_list_from_col(top_list, lview); + populate_from_col(top_list->children[0], lview.child()); + NANOARROW_THROW_NOT_OK( + ArrowArrayFinishBuilding(input_array.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); + + // slice at the fixed-size-list level only. the synthesized offsets must stay absolute, + // so that the outer column_view offset indexes them correctly + input_array->length = 2; + top_list->offset = 1; + top_list->length = 2; + + ArrowDeviceArray input_device_array; + input_device_array.device_id = rmm::get_current_cuda_device().value(); + input_device_array.device_type = ARROW_DEVICE_CUDA; + input_device_array.sync_event = nullptr; + memcpy(&input_device_array.array, input_array.get(), sizeof(ArrowArray)); + + auto got_cudf_table_view = cudf::from_arrow_device(input_schema.get(), &input_device_array); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected_table_view, *got_cudf_table_view); +} + TEST_F(FromArrowDeviceTest, StructColumn) { using vector_of_columns = std::vector>; diff --git a/cpp/tests/interop/from_arrow_host_test.cpp b/cpp/tests/interop/from_arrow_host_test.cpp index 3cb451165f79..4eba65ee3541 100644 --- a/cpp/tests/interop/from_arrow_host_test.cpp +++ b/cpp/tests/interop/from_arrow_host_test.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -526,6 +527,186 @@ TEST_F(FromArrowHostDeviceTest, NestedList) CUDF_TEST_EXPECT_TABLES_EQUAL(got_cudf_table->view(), from_struct); } +namespace { + +// Build a struct schema carrying a single fixed_size_list[width] child named "a". +// ArrowSchemaInitFromType cannot be used for NANOARROW_TYPE_FIXED_SIZE_LIST: there is no +// unambiguous format template for it, so it fails with EINVAL. ArrowSchemaSetTypeFixedSize +// is the supported path, and it allocates the "item" child with a NULL format, so the child +// type still has to be set explicitly. +nanoarrow::UniqueSchema make_fixed_size_list_schema(int32_t width, bool nullable) +{ + nanoarrow::UniqueSchema schema; + ArrowSchemaInit(schema.get()); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetTypeStruct(schema.get(), 1)); + + NANOARROW_THROW_NOT_OK( + ArrowSchemaSetTypeFixedSize(schema->children[0], NANOARROW_TYPE_FIXED_SIZE_LIST, width)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0], "a")); + schema->children[0]->flags = nullable ? ARROW_FLAG_NULLABLE : 0; + + NANOARROW_THROW_NOT_OK( + ArrowSchemaSetType(schema->children[0]->children[0], NANOARROW_TYPE_INT64)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0]->children[0], "element")); + schema->children[0]->children[0]->flags = 0; + + return schema; +} + +// Build the matching ArrowArray. `values` holds num_rows * width int64 child elements. +// A fixed_size_list array has no offsets buffer, only validity, so nothing is written to +// buffer index 1 of the list level. +nanoarrow::UniqueArray make_fixed_size_list_array(ArrowSchema* schema, + std::vector const& values, + int64_t num_rows, + std::vector const& list_validity = {}) +{ + nanoarrow::UniqueArray array; + NANOARROW_THROW_NOT_OK(ArrowArrayInitFromSchema(array.get(), schema, nullptr)); + array->length = num_rows; + array->null_count = 0; + + auto* list_array = array->children[0]; + list_array->length = num_rows; + list_array->null_count = 0; + if (!list_validity.empty()) { + ArrowBitmap bitmap; + ArrowBitmapInit(&bitmap); + NANOARROW_THROW_NOT_OK(ArrowBitmapReserve(&bitmap, list_validity.size())); + ArrowBitmapAppendInt8Unsafe( + &bitmap, reinterpret_cast(list_validity.data()), list_validity.size()); + ArrowArraySetValidityBitmap(list_array, &bitmap); + list_array->null_count = + num_rows - + ArrowBitCountSet(ArrowArrayValidityBitmap(list_array)->buffer.data, 0, list_validity.size()); + } + + auto* values_array = list_array->children[0]; + NANOARROW_THROW_NOT_OK(ArrowBufferAppend(ArrowArrayBuffer(values_array, 1), + reinterpret_cast(values.data()), + values.size() * sizeof(int64_t))); + values_array->length = values.size(); + values_array->null_count = 0; + + NANOARROW_THROW_NOT_OK( + ArrowArrayFinishBuilding(array.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); + return array; +} + +ArrowDeviceArray as_host_device_array(nanoarrow::UniqueArray const& array) +{ + ArrowDeviceArray input; + memcpy(&input.array, array.get(), sizeof(ArrowArray)); + input.device_id = -1; + input.device_type = ARROW_DEVICE_CPU; + return input; +} + +} // namespace + +TEST_F(FromArrowHostDeviceTest, FixedSizeListColumn) +{ + constexpr int32_t width = 3; + constexpr int64_t num_rows = 4; + std::vector values{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + + auto expected_col = + cudf::test::lists_column_wrapper{{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; + cudf::table_view expected_table_view({expected_col}); + + auto input_schema = make_fixed_size_list_schema(width, /*nullable=*/false); + auto input_array = make_fixed_size_list_array(input_schema.get(), values, num_rows); + auto input = as_host_device_array(input_array); + + auto got_cudf_table = cudf::from_arrow_host(input_schema.get(), &input); + EXPECT_EQ(got_cudf_table->get_column(0).type(), cudf::data_type{cudf::type_id::LIST}); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table_view, got_cudf_table->view()); + + ArrowDeviceArray direct_input; + memcpy(&direct_input.array, input_array->children[0], sizeof(ArrowArray)); + direct_input.device_id = -1; + direct_input.device_type = ARROW_DEVICE_CPU; + direct_input.sync_event = nullptr; + auto got_direct_col = cudf::from_arrow_host_column(input_schema->children[0], &direct_input); + EXPECT_EQ(got_direct_col->type(), cudf::data_type{cudf::type_id::LIST}); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_col, got_direct_col->view()); + + auto got_cudf_col = cudf::from_arrow_host_column(input_schema.get(), &input); + EXPECT_EQ(got_cudf_col->type(), cudf::data_type{cudf::type_id::STRUCT}); + auto got_cudf_col_view = got_cudf_col->view(); + cudf::table_view from_struct{ + std::vector(got_cudf_col_view.child_begin(), got_cudf_col_view.child_end())}; + CUDF_TEST_EXPECT_TABLES_EQUAL(got_cudf_table->view(), from_struct); +} + +TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnNulls) +{ + constexpr int32_t width = 2; + constexpr int64_t num_rows = 4; + // a null fixed-size-list row still occupies `width` child slots, so the child data is + // dense and the offsets stay exact multiples of the width + std::vector values{1, 2, 3, 4, 5, 6, 7, 8}; + std::vector list_validity{1, 0, 1, 0}; + + // lists_column_wrapper cannot express this: it encodes a null row as a repeated offset + // and drops that row's child values, which breaks the multiple-of-width invariant. + auto child = + cudf::test::fixed_width_column_wrapper(values.begin(), values.end()).release(); + auto offsets = cudf::test::fixed_width_column_wrapper{0, 2, 4, 6, 8}.release(); + auto [null_mask, null_count] = + cudf::test::detail::make_null_mask(list_validity.begin(), list_validity.end()); + auto expected_col = cudf::make_lists_column( + num_rows, std::move(offsets), std::move(child), null_count, std::move(null_mask)); + cudf::table_view expected_table_view({expected_col->view()}); + + auto input_schema = make_fixed_size_list_schema(width, /*nullable=*/true); + auto input_array = + make_fixed_size_list_array(input_schema.get(), values, num_rows, list_validity); + auto input = as_host_device_array(input_array); + + auto got_cudf_table = cudf::from_arrow_host(input_schema.get(), &input); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table_view, got_cudf_table->view()); +} + +TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnSliced) +{ + constexpr int32_t width = 3; + constexpr int64_t num_rows = 4; + std::vector values{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + + auto full_col = + cudf::test::lists_column_wrapper{{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; + auto sliced = cudf::slice(full_col, {1, 3}); + cudf::table_view expected_table_view({sliced.front()}); + + auto input_schema = make_fixed_size_list_schema(width, /*nullable=*/false); + auto input_array = make_fixed_size_list_array(input_schema.get(), values, num_rows); + auto input = as_host_device_array(input_array); + // this is what catches an incorrect `input->offset * width`, since the child range must + // start at row 1 * width rather than at zero + slice_host_nanoarrow(&input.array, 1, 3); + + auto got_cudf_table = cudf::from_arrow_host(input_schema.get(), &input); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table_view, got_cudf_table->view()); +} + +TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnZeroLength) +{ + constexpr int32_t width = 3; + + auto expected_col = cudf::test::lists_column_wrapper{}; + cudf::table_view expected_table_view({expected_col}); + + auto input_schema = make_fixed_size_list_schema(width, /*nullable=*/false); + auto input_array = make_fixed_size_list_array(input_schema.get(), {}, 0); + auto input = as_host_device_array(input_array); + + auto got_cudf_table = cudf::from_arrow_host(input_schema.get(), &input); + EXPECT_EQ(got_cudf_table->num_rows(), 0); + EXPECT_EQ(got_cudf_table->get_column(0).type(), cudf::data_type{cudf::type_id::LIST}); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table_view, got_cudf_table->view()); +} + TEST_F(FromArrowHostDeviceTest, StructColumn) { // Create cudf table diff --git a/cpp/tests/interop/from_arrow_stream_test.cpp b/cpp/tests/interop/from_arrow_stream_test.cpp index 74f9fc1df31d..c9983f758f2e 100644 --- a/cpp/tests/interop/from_arrow_stream_test.cpp +++ b/cpp/tests/interop/from_arrow_stream_test.cpp @@ -1,10 +1,11 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include #include +#include #include #include @@ -118,3 +119,91 @@ TEST_F(FromArrowStreamTest, EmptyChunkedTest) auto result = cudf::from_arrow_stream_column(&stream); CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view(), expected->view()); } + +namespace { + +// Builds a struct schema with one fixed_size_list[width] child. ArrowSchemaInitFromType +// does not support NANOARROW_TYPE_FIXED_SIZE_LIST, so ArrowSchemaSetTypeFixedSize is used and +// the allocated "item" child gets its type set explicitly. +nanoarrow::UniqueSchema make_fixed_size_list_stream_schema(int32_t width) +{ + nanoarrow::UniqueSchema schema; + ArrowSchemaInit(schema.get()); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetTypeStruct(schema.get(), 1)); + + NANOARROW_THROW_NOT_OK( + ArrowSchemaSetTypeFixedSize(schema->children[0], NANOARROW_TYPE_FIXED_SIZE_LIST, width)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0], "a")); + schema->children[0]->flags = 0; + + NANOARROW_THROW_NOT_OK( + ArrowSchemaSetType(schema->children[0]->children[0], NANOARROW_TYPE_INT64)); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0]->children[0], "element")); + schema->children[0]->children[0]->flags = 0; + + return schema; +} + +nanoarrow::UniqueArray make_fixed_size_list_chunk(ArrowSchema* schema, + std::vector const& values, + int64_t num_rows) +{ + nanoarrow::UniqueArray array; + NANOARROW_THROW_NOT_OK(ArrowArrayInitFromSchema(array.get(), schema, nullptr)); + array->length = num_rows; + array->null_count = 0; + + auto* list_array = array->children[0]; + list_array->length = num_rows; + list_array->null_count = 0; + + auto* values_array = list_array->children[0]; + NANOARROW_THROW_NOT_OK(ArrowBufferAppend(ArrowArrayBuffer(values_array, 1), + reinterpret_cast(values.data()), + values.size() * sizeof(int64_t))); + values_array->length = values.size(); + values_array->null_count = 0; + + NANOARROW_THROW_NOT_OK( + ArrowArrayFinishBuilding(array.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); + return array; +} + +} // namespace + +// exercises make_empty_column_from_schema, which builds the column from the schema alone +TEST_F(FromArrowStreamTest, FixedSizeListEmptyTest) +{ + auto schema = make_fixed_size_list_stream_schema(3); + + ArrowArrayStream stream; + makeStreamFromArrays({}, std::move(schema), &stream); + + auto result = cudf::from_arrow_stream(&stream); + EXPECT_EQ(result->num_rows(), 0); + EXPECT_EQ(result->get_column(0).type(), cudf::data_type{cudf::type_id::LIST}); +} + +// exercises concatenate over columns whose offsets were synthesized rather than copied +TEST_F(FromArrowStreamTest, FixedSizeListChunkedTest) +{ + constexpr int32_t width = 2; + auto schema = make_fixed_size_list_stream_schema(width); + + std::vector arrays; + for (auto i = 0; i < 3; ++i) { + auto base = static_cast(i * 4); + arrays.push_back(make_fixed_size_list_chunk( + schema.get(), {base + 1, base + 2, base + 3, base + 4}, /*num_rows=*/2)); + } + + auto expected_col = + cudf::test::lists_column_wrapper{{1, 2}, {3, 4}, {5, 6}, {7, 8}, {9, 10}, {11, 12}}; + cudf::table_view expected_table_view({expected_col}); + + ArrowArrayStream stream; + makeStreamFromArrays(std::move(arrays), std::move(schema), &stream); + + auto result = cudf::from_arrow_stream(&stream); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table_view, result->view()); +} diff --git a/python/pylibcudf/tests/test_table.py b/python/pylibcudf/tests/test_table.py index 91ff25055565..65c1c47186a6 100644 --- a/python/pylibcudf/tests/test_table.py +++ b/python/pylibcudf/tests/test_table.py @@ -107,6 +107,19 @@ def test_from_arrow_zero_column_preserves_num_rows(): assert tbl.num_rows() == 5 +def test_from_arrow_fixed_size_list_normalizes_to_list(): + fixed = pa.array( + [[1, 2, 3], [4, 5, 6]], type=pa.list_(pa.int64(), list_size=3) + ) + expected = pa.array([[1, 2, 3], [4, 5, 6]], type=pa.list_(pa.int64())) + + column = plc.Column.from_arrow(fixed) + assert column.to_arrow().equals(expected) + + table = plc.Table.from_arrow(pa.table({"a": fixed})) + assert_table_eq(table.to_arrow(), pa.table({"a": expected})) + + def test_to_arrow_zero_column_preserves_num_rows(): arrow = plc.Table([], num_rows=5).to_arrow() assert arrow.num_columns == 0 From 4dcbca84d85e7a44ebb97ed5d9ffb6986f65af58 Mon Sep 17 00:00:00 2001 From: Oleksii Gubanov <153339237+0guban0v@users.noreply.github.com> Date: Fri, 7 Aug 2026 00:41:18 -0700 Subject: [PATCH 2/3] Fix fixed-size-list null test comparison --- cpp/tests/interop/from_arrow_host_test.cpp | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/cpp/tests/interop/from_arrow_host_test.cpp b/cpp/tests/interop/from_arrow_host_test.cpp index 4eba65ee3541..c09f756d2921 100644 --- a/cpp/tests/interop/from_arrow_host_test.cpp +++ b/cpp/tests/interop/from_arrow_host_test.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -657,7 +658,6 @@ TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnNulls) cudf::test::detail::make_null_mask(list_validity.begin(), list_validity.end()); auto expected_col = cudf::make_lists_column( num_rows, std::move(offsets), std::move(child), null_count, std::move(null_mask)); - cudf::table_view expected_table_view({expected_col->view()}); auto input_schema = make_fixed_size_list_schema(width, /*nullable=*/true); auto input_array = @@ -665,7 +665,16 @@ TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnNulls) auto input = as_host_device_array(input_array); auto got_cudf_table = cudf::from_arrow_host(input_schema.get(), &input); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table_view, got_cudf_table->view()); + auto const expected_lists = cudf::lists_column_view(expected_col->view()); + auto const got_lists = cudf::lists_column_view(got_cudf_table->get_column(0)); + + EXPECT_TRUE(cudf::has_nonempty_nulls(got_lists.parent())); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_lists.offsets(), got_lists.offsets()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_lists.child(), got_lists.child()); + + auto const expected_logical = cudf::purge_nonempty_nulls(expected_lists.parent()); + auto const got_logical = cudf::purge_nonempty_nulls(got_lists.parent()); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(*expected_logical, *got_logical); } TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnSliced) From 8599f6f2b6fcb2fde9f9fe7b61dbfb71fdb374c2 Mon Sep 17 00:00:00 2001 From: Oleksii Gubanov <153339237+0guban0v@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:21:27 -0700 Subject: [PATCH 3/3] Address fixed-size-list ingress review feedback --- cpp/src/interop/arrow_utilities.cpp | 41 +++++ cpp/src/interop/arrow_utilities.hpp | 28 ++++ cpp/src/interop/from_arrow_device.cu | 32 ++-- cpp/src/interop/from_arrow_host.cu | 46 +++--- cpp/tests/interop/from_arrow_device_test.cpp | 161 ++++++++++++++++++- cpp/tests/interop/from_arrow_host_test.cpp | 94 ++++++++++- cpp/tests/interop/from_arrow_stream_test.cpp | 105 +++++++++++- python/pylibcudf/tests/test_table.py | 28 +++- 8 files changed, 488 insertions(+), 47 deletions(-) diff --git a/cpp/src/interop/arrow_utilities.cpp b/cpp/src/interop/arrow_utilities.cpp index ec7bab8d6aae..5d8ea4784eaf 100644 --- a/cpp/src/interop/arrow_utilities.cpp +++ b/cpp/src/interop/arrow_utilities.cpp @@ -11,6 +11,7 @@ #include +#include #include namespace cudf { @@ -77,9 +78,49 @@ size_type fixed_size_list_width(ArrowSchemaView const* arrow_view) CUDF_EXPECTS(arrow_view->fixed_size >= 0, "fixed-size-list width must be non-negative", std::invalid_argument); + CUDF_EXPECTS(arrow_view->fixed_size <= std::numeric_limits::max(), + "fixed-size-list width exceeds cuDF's maximum supported row count (cudf::size_type)", + std::overflow_error); return static_cast(arrow_view->fixed_size); } +fixed_size_list_layout get_fixed_size_list_layout(ArrowSchemaView const* arrow_view, + ArrowArray const* input) +{ + CUDF_EXPECTS(input->offset >= 0 && input->length >= 0, + "fixed-size-list offset and length must be non-negative", + std::invalid_argument); + + constexpr auto max_size = static_cast(std::numeric_limits::max()); + CUDF_EXPECTS(input->length < max_size, + "fixed-size-list length exceeds cuDF's maximum supported row count " + "(cudf::size_type)", + std::overflow_error); + CUDF_EXPECTS(input->offset <= std::numeric_limits::max() - input->length, + "fixed-size-list row bounds overflow Arrow's int64 representation", + std::overflow_error); + + auto const width = fixed_size_list_width(arrow_view); + auto const num_rows = static_cast(input->length); + auto const row_end = input->offset + input->length; + + // Width zero is valid for a foreign Arrow producer even though nanoarrow's schema builder + // rejects it. Its offsets and child bounds are all zero. + if (width == 0) { return {width, num_rows, input->offset, row_end, 0, 0, 0}; } + + CUDF_EXPECTS(row_end <= std::numeric_limits::max() / width, + "fixed-size-list child bounds overflow Arrow's int64 representation", + std::overflow_error); + auto const child_length = input->length * width; + CUDF_EXPECTS(child_length <= max_size, + "Number of fixed-size-list child elements exceeds cuDF's maximum supported " + "row count (cudf::size_type)", + std::overflow_error); + + return { + width, num_rows, input->offset, row_end, input->offset * width, child_length, row_end * width}; +} + ArrowType id_to_arrow_type(cudf::type_id id) { switch (id) { diff --git a/cpp/src/interop/arrow_utilities.hpp b/cpp/src/interop/arrow_utilities.hpp index 93c65761796b..58958ae046be 100644 --- a/cpp/src/interop/arrow_utilities.hpp +++ b/cpp/src/interop/arrow_utilities.hpp @@ -36,17 +36,45 @@ data_type arrow_to_cudf_type(ArrowSchemaView const* arrow_view); */ bool is_fixed_size_list(ArrowSchemaView const* arrow_view); +/** + * @brief Validated physical bounds for an Arrow fixed-size-list array + */ +struct fixed_size_list_layout { + size_type width; ///< Child elements per row + size_type num_rows; ///< Number of output rows + int64_t row_offset; ///< First logical row in the Arrow array + int64_t row_end; ///< One-past-last logical row + int64_t child_offset; ///< First referenced child element + int64_t child_length; ///< Number of referenced child elements + int64_t child_end; ///< One-past-last referenced child element +}; + /** * @brief Return the number of child elements per row of a fixed-size-list schema * * @throw cudf::data_type_error if `arrow_view` is not a fixed-size-list * @throw std::invalid_argument if the declared width is negative + * @throw std::overflow_error if the declared width exceeds `size_type` * * @param arrow_view SchemaView to pull the fixed size from * @return Number of child elements per row */ size_type fixed_size_list_width(ArrowSchemaView const* arrow_view); +/** + * @brief Validate and compute fixed-size-list row and child bounds + * + * @throw std::invalid_argument if row metadata is negative + * @throw std::overflow_error if Arrow bounds overflow `int64_t` or output lengths exceed + * `size_type` + * + * @param arrow_view Fixed-size-list schema view + * @param input Arrow array carrying row offset and length + * @return Validated source bounds and output sizes + */ +fixed_size_list_layout get_fixed_size_list_layout(ArrowSchemaView const* arrow_view, + ArrowArray const* input); + /** * @brief Map cudf column type id to ArrowType id * diff --git a/cpp/src/interop/from_arrow_device.cu b/cpp/src/interop/from_arrow_device.cu index af65503b9b78..fcb893afc083 100644 --- a/cpp/src/interop/from_arrow_device.cu +++ b/cpp/src/interop/from_arrow_device.cu @@ -322,8 +322,22 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()( CUDF_EXPECTS(schema->type != NANOARROW_TYPE_LARGE_LIST, "Large list types are not supported", cudf::data_type_error); - size_type const num_rows = input->length; - size_type const offset = input->offset; + auto const fixed_size = is_fixed_size_list(schema); + auto const layout = + fixed_size ? get_fixed_size_list_layout(schema, input) : fixed_size_list_layout{}; + + if (fixed_size) { + constexpr auto max_size = static_cast(std::numeric_limits::max()); + CUDF_EXPECTS(layout.row_end < max_size && layout.child_end <= max_size, + "fixed-size-list device bounds exceed cuDF's maximum supported row count " + "(cudf::size_type)", + std::overflow_error); + CUDF_EXPECTS(input->children[0]->length >= layout.child_end, + "fixed-size-list child is shorter than its parent layout requires", + std::invalid_argument); + } + size_type const num_rows = fixed_size ? layout.num_rows : input->length; + size_type const offset = fixed_size ? static_cast(layout.row_offset) : input->offset; size_type const null_count = input->null_count; ArrowSchemaView child_schema_view; @@ -338,24 +352,16 @@ dispatch_tuple_t dispatch_from_arrow_device::operator()( // so that when `get_sliced_child` is called, we still produce the right result column_view offsets_view; size_type max_child_offset = 0; - if (is_fixed_size_list(schema)) { + if (fixed_size) { // fixed-size-list arrays carry no offsets buffer, so synthesize {0, w, 2w, ...}. // these are absolute rather than normalized because the outer column_view applies a // single offset to both the null mask and the children. - auto const width = static_cast(fixed_size_list_width(schema)); - auto const num_offsets = static_cast(offset) + num_rows + 1; - auto const child_end = (num_offsets - 1) * width; - constexpr auto max_size = static_cast(std::numeric_limits::max()); - CUDF_EXPECTS(num_offsets <= max_size && child_end <= max_size, - "fixed-size-list offsets exceed cuDF's maximum supported row count " - "(cudf::size_type).", - std::overflow_error); - max_child_offset = (num_rows == 0) ? 0 : static_cast(child_end); + max_child_offset = (num_rows == 0) ? 0 : static_cast(layout.child_end); if (num_rows == 0) { offsets_view = column_view{data_type{type_id::INT32}, 0, nullptr, nullptr, 0, 0}; } else { owned.emplace_back(make_fixed_size_list_offsets( - static_cast(num_offsets), static_cast(width), stream, mr)); + static_cast(layout.row_end) + 1, layout.width, stream, mr)); offsets_view = owned.back()->view(); } } else { diff --git a/cpp/src/interop/from_arrow_host.cu b/cpp/src/interop/from_arrow_host.cu index 89b57ecb1e0a..81a73b6b3ebb 100644 --- a/cpp/src/interop/from_arrow_host.cu +++ b/cpp/src/interop/from_arrow_host.cu @@ -300,33 +300,28 @@ std::tuple, int64_t, int64_t> get_fixed_size_list_offset rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - auto const width = static_cast(fixed_size_list_width(schema)); - auto const offset = input->offset * width; - auto const length = input->length * width; - - constexpr auto max_size = static_cast(std::numeric_limits::max()); - CUDF_EXPECTS(length <= max_size, - "Number of fixed-size-list child elements exceeds cuDF's maximum supported " - "row count (cudf::size_type).", - std::overflow_error); - - return std::tuple{ - make_fixed_size_list_offsets( - static_cast(input->length) + 1, static_cast(width), stream, mr), - offset, - length}; + auto const layout = get_fixed_size_list_layout(schema, input); + CUDF_EXPECTS(input->children[0]->length >= layout.child_end, + "fixed-size-list child is shorter than its parent layout requires", + std::invalid_argument); + + return std::tuple{make_fixed_size_list_offsets(layout.num_rows + 1, layout.width, stream, mr), + layout.child_offset, + layout.child_length}; } template <> std::unique_ptr dispatch_copy_from_arrow_host::operator()( ArrowSchemaView const* schema, ArrowArray const* input, data_type type, bool skip_mask) { + CUDF_EXPECTS(input->length >= 0, "Number of rows must be non-negative.", std::invalid_argument); CUDF_EXPECTS( - input->length + 1 <= static_cast(std::numeric_limits::max()), + input->length < static_cast(std::numeric_limits::max()), "Number of rows exceeds cuDF's maximum supported row count (cudf::size_type).", std::overflow_error); - auto [offsets_column, offset, length] = is_fixed_size_list(schema) + auto const fixed_size = is_fixed_size_list(schema); + auto [offsets_column, offset, length] = fixed_size ? get_fixed_size_list_offsets(schema, input, stream, mr) : get_offsets_column(schema, input, stream, mr); @@ -335,8 +330,16 @@ std::unique_ptr dispatch_copy_from_arrow_host::operator()children[0]); + if (fixed_size) { + CUDF_EXPECTS(child_array.offset >= 0, + "fixed-size-list child offset must be non-negative", + std::invalid_argument); + CUDF_EXPECTS(offset <= std::numeric_limits::max() - child_array.offset, + "fixed-size-list child offset overflows Arrow's int64 representation", + std::overflow_error); + } child_array.offset += offset; - child_array.length = std::min(length, child_array.length); + child_array.length = fixed_size ? length : std::min(length, child_array.length); auto child_column = get_column_copy(&view, &child_array, child_type, skip_mask, stream, mr); @@ -437,11 +440,8 @@ std::unique_ptr make_fixed_size_list_offsets(size_type size, auto offsets = make_numeric_column(data_type{type_id::INT32}, size, mask_state::UNALLOCATED, stream, mr); auto d_offsets = offsets->mutable_view().begin(); - thrust::sequence(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - d_offsets, - d_offsets + size, - size_type{0}, - width); + thrust::sequence( + rmm::exec_policy_nosync(stream, mr), d_offsets, d_offsets + size, size_type{0}, width); return offsets; } diff --git a/cpp/tests/interop/from_arrow_device_test.cpp b/cpp/tests/interop/from_arrow_device_test.cpp index eb8f899fa089..3c026f412451 100644 --- a/cpp/tests/interop/from_arrow_device_test.cpp +++ b/cpp/tests/interop/from_arrow_device_test.cpp @@ -19,12 +19,17 @@ #include #include #include +#include #include #include #include #include +#include +#include +#include + struct FromArrowDeviceTest : public cudf::test::BaseFixture {}; template @@ -286,7 +291,7 @@ void populate_fixed_size_list_from_col(ArrowArray* arr, cudf::lists_column_view // ArrowSchemaInitFromType does not support NANOARROW_TYPE_FIXED_SIZE_LIST (no format // template, returns EINVAL); ArrowSchemaSetTypeFixedSize is the supported path and leaves // the allocated "item" child with a NULL format, so the child type is set explicitly. -nanoarrow::UniqueSchema make_fixed_size_list_device_schema(int32_t width) +nanoarrow::UniqueSchema make_fixed_size_list_device_schema(int32_t width, bool nullable = false) { nanoarrow::UniqueSchema schema; ArrowSchemaInit(schema.get()); @@ -295,7 +300,7 @@ nanoarrow::UniqueSchema make_fixed_size_list_device_schema(int32_t width) NANOARROW_THROW_NOT_OK( ArrowSchemaSetTypeFixedSize(schema->children[0], NANOARROW_TYPE_FIXED_SIZE_LIST, width)); NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0], "a")); - schema->children[0]->flags = 0; + schema->children[0]->flags = nullable ? ARROW_FLAG_NULLABLE : 0; NANOARROW_THROW_NOT_OK( ArrowSchemaSetType(schema->children[0]->children[0], NANOARROW_TYPE_INT64)); @@ -343,6 +348,7 @@ TEST_F(FromArrowDeviceTest, FixedSizeListColumn) auto got_direct_col = cudf::from_arrow_device_column(input_schema->children[0], &direct_device_array); EXPECT_EQ(got_direct_col->type(), cudf::data_type{cudf::type_id::LIST}); + EXPECT_EQ(got_direct_col.get_deleter().owned_mem_.size(), 1); CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_table_view.column(0), *got_direct_col); auto got_cudf_col = cudf::from_arrow_device_column(input_schema.get(), &input_device_array); @@ -387,6 +393,157 @@ TEST_F(FromArrowDeviceTest, FixedSizeListColumnSliced) CUDF_TEST_EXPECT_TABLES_EQUAL(expected_table_view, *got_cudf_table_view); } +TEST_F(FromArrowDeviceTest, FixedSizeListColumnEmpty) +{ + auto expected = cudf::test::lists_column_wrapper{}; + auto input_schema = make_fixed_size_list_device_schema(3); + + nanoarrow::UniqueArray input_array; + NANOARROW_THROW_NOT_OK(ArrowArrayInitFromSchema(input_array.get(), input_schema.get(), nullptr)); + auto* list_array = input_array->children[0]; + cudf::lists_column_view expected_view{expected}; + populate_fixed_size_list_from_col(list_array, expected_view); + populate_from_col(list_array->children[0], expected_view.child()); + NANOARROW_THROW_NOT_OK( + ArrowArrayFinishBuilding(input_array.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); + + ArrowDeviceArray input; + memcpy(&input.array, list_array, sizeof(ArrowArray)); + input.device_id = rmm::get_current_cuda_device().value(); + input.device_type = ARROW_DEVICE_CUDA; + input.sync_event = nullptr; + + auto result = cudf::from_arrow_device_column(input_schema->children[0], &input); + EXPECT_EQ(result->type(), cudf::data_type{cudf::type_id::LIST}); + EXPECT_EQ(result->size(), 0); + EXPECT_TRUE(result.get_deleter().owned_mem_.empty()); +} + +TEST_F(FromArrowDeviceTest, FixedSizeListColumnNulls) +{ + constexpr cudf::size_type num_rows = 4; + std::vector validity{1, 0, 1, 0}; + auto child = cudf::test::fixed_width_column_wrapper{1, 2, 3, 4, 5, 6, 7, 8}.release(); + auto offsets = cudf::test::fixed_width_column_wrapper{0, 2, 4, 6, 8}.release(); + auto [null_mask, null_count] = + cudf::test::detail::make_null_mask(validity.begin(), validity.end()); + auto expected = cudf::make_lists_column( + num_rows, std::move(offsets), std::move(child), null_count, std::move(null_mask)); + + auto input_schema = make_fixed_size_list_device_schema(2, /*nullable=*/true); + nanoarrow::UniqueArray input_array; + NANOARROW_THROW_NOT_OK(ArrowArrayInitFromSchema(input_array.get(), input_schema.get(), nullptr)); + auto* list_array = input_array->children[0]; + cudf::lists_column_view expected_view{*expected}; + populate_fixed_size_list_from_col(list_array, expected_view); + populate_from_col(list_array->children[0], expected_view.child()); + NANOARROW_THROW_NOT_OK( + ArrowArrayFinishBuilding(input_array.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); + + ArrowDeviceArray input; + memcpy(&input.array, list_array, sizeof(ArrowArray)); + input.device_id = rmm::get_current_cuda_device().value(); + input.device_type = ARROW_DEVICE_CUDA; + input.sync_event = nullptr; + + auto result = cudf::from_arrow_device_column(input_schema->children[0], &input); + auto result_lists = cudf::lists_column_view{*result}; + EXPECT_EQ(result.get_deleter().owned_mem_.size(), 1); + EXPECT_TRUE(cudf::has_nonempty_nulls(result_lists.parent())); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_view.offsets(), result_lists.offsets()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_view.child(), result_lists.child()); + + auto expected_logical = cudf::purge_nonempty_nulls(expected_view.parent()); + auto result_logical = cudf::purge_nonempty_nulls(result_lists.parent()); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(*expected_logical, *result_logical); +} + +TEST_F(FromArrowDeviceTest, FixedSizeListColumnLarge) +{ + constexpr cudf::size_type width = 2; + + for (auto const num_rows : {cudf::size_type{1024}, cudf::size_type{1025}}) { + SCOPED_TRACE(num_rows); + std::vector values(num_rows * width); + std::iota(values.begin(), values.end(), int64_t{0}); + std::vector offsets(num_rows + 1); + for (cudf::size_type i = 0; i <= num_rows; ++i) { + offsets[i] = i * width; + } + + auto child = + cudf::test::fixed_width_column_wrapper(values.begin(), values.end()).release(); + auto offsets_col = + cudf::test::fixed_width_column_wrapper(offsets.begin(), offsets.end()) + .release(); + auto expected = cudf::make_lists_column( + num_rows, std::move(offsets_col), std::move(child), 0, rmm::device_buffer{}); + + auto input_schema = make_fixed_size_list_device_schema(width); + nanoarrow::UniqueArray input_array; + NANOARROW_THROW_NOT_OK( + ArrowArrayInitFromSchema(input_array.get(), input_schema.get(), nullptr)); + auto* list_array = input_array->children[0]; + cudf::lists_column_view expected_view{*expected}; + populate_fixed_size_list_from_col(list_array, expected_view); + populate_from_col(list_array->children[0], expected_view.child()); + NANOARROW_THROW_NOT_OK( + ArrowArrayFinishBuilding(input_array.get(), NANOARROW_VALIDATION_LEVEL_NONE, nullptr)); + + ArrowDeviceArray input; + memcpy(&input.array, list_array, sizeof(ArrowArray)); + input.device_id = rmm::get_current_cuda_device().value(); + input.device_type = ARROW_DEVICE_CUDA; + input.sync_event = nullptr; + + auto result = cudf::from_arrow_device_column(input_schema->children[0], &input); + EXPECT_EQ(result.get_deleter().owned_mem_.size(), 1); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected->view(), *result); + } +} + +TEST_F(FromArrowDeviceTest, FixedSizeListInvalidBounds) +{ + auto input_schema = make_fixed_size_list_device_schema(3); + nanoarrow::UniqueArray input_array; + NANOARROW_THROW_NOT_OK(ArrowArrayInitFromSchema(input_array.get(), input_schema.get(), nullptr)); + + ArrowDeviceArray input; + memcpy(&input.array, input_array->children[0], sizeof(ArrowArray)); + input.device_id = rmm::get_current_cuda_device().value(); + input.device_type = ARROW_DEVICE_CUDA; + input.sync_event = nullptr; + + input.array.offset = -1; + EXPECT_THROW(cudf::from_arrow_device_column(input_schema->children[0], &input), + std::invalid_argument); + + input.array.offset = 0; + input.array.length = -1; + EXPECT_THROW(cudf::from_arrow_device_column(input_schema->children[0], &input), + std::invalid_argument); + + input.array.length = std::numeric_limits::max(); + EXPECT_THROW(cudf::from_arrow_device_column(input_schema->children[0], &input), + std::overflow_error); + + input.array.offset = std::numeric_limits::max(); + input.array.length = 1; + EXPECT_THROW(cudf::from_arrow_device_column(input_schema->children[0], &input), + std::overflow_error); + + input.array.offset = std::numeric_limits::max() / 2; + input.array.length = 1; + EXPECT_THROW(cudf::from_arrow_device_column(input_schema->children[0], &input), + std::overflow_error); + + input.array.offset = 0; + input.array.length = 1; + input.array.children[0]->length = 2; + EXPECT_THROW(cudf::from_arrow_device_column(input_schema->children[0], &input), + std::invalid_argument); +} + TEST_F(FromArrowDeviceTest, StructColumn) { using vector_of_columns = std::vector>; diff --git a/cpp/tests/interop/from_arrow_host_test.cpp b/cpp/tests/interop/from_arrow_host_test.cpp index c09f756d2921..f6c1634f8a9f 100644 --- a/cpp/tests/interop/from_arrow_host_test.cpp +++ b/cpp/tests/interop/from_arrow_host_test.cpp @@ -26,6 +26,10 @@ #include +#include +#include +#include + // create a cudf::table and equivalent arrow table with host memory std::tuple, nanoarrow::UniqueSchema, nanoarrow::UniqueArray> get_nanoarrow_host_tables(cudf::size_type length) @@ -664,7 +668,7 @@ TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnNulls) make_fixed_size_list_array(input_schema.get(), values, num_rows, list_validity); auto input = as_host_device_array(input_array); - auto got_cudf_table = cudf::from_arrow_host(input_schema.get(), &input); + auto got_cudf_table = cudf::from_arrow_host(input_schema.get(), &input); auto const expected_lists = cudf::lists_column_view(expected_col->view()); auto const got_lists = cudf::lists_column_view(got_cudf_table->get_column(0)); @@ -716,6 +720,94 @@ TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnZeroLength) CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table_view, got_cudf_table->view()); } +TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnZeroWidth) +{ + constexpr cudf::size_type num_rows = 3; + auto offsets = cudf::test::fixed_width_column_wrapper{0, 0, 0, 0}.release(); + auto child = cudf::test::fixed_width_column_wrapper{}.release(); + auto expected = cudf::make_lists_column(num_rows, std::move(offsets), std::move(child), 0, {}); + + // nanoarrow's schema builder rejects width zero, but ArrowSchemaView accepts it from a + // foreign producer. Replace a normally constructed fixed-size-list format to exercise it. + auto input_schema = make_fixed_size_list_schema(1, /*nullable=*/false); + NANOARROW_THROW_NOT_OK(ArrowSchemaSetFormat(input_schema->children[0], "+w:0")); + auto input_array = make_fixed_size_list_array(input_schema.get(), {}, num_rows); + auto input = as_host_device_array(input_array); + + auto result = cudf::from_arrow_host(input_schema.get(), &input); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected->view(), result->get_column(0)); +} + +TEST_F(FromArrowHostDeviceTest, FixedSizeListColumnLarge) +{ + constexpr cudf::size_type width = 2; + + for (auto const num_rows : {cudf::size_type{1024}, cudf::size_type{1025}}) { + SCOPED_TRACE(num_rows); + std::vector values(num_rows * width); + std::iota(values.begin(), values.end(), int64_t{0}); + std::vector offsets(num_rows + 1); + for (cudf::size_type i = 0; i <= num_rows; ++i) { + offsets[i] = i * width; + } + + auto expected_offsets = + cudf::test::fixed_width_column_wrapper(offsets.begin(), offsets.end()); + auto expected_child = + cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); + + auto input_schema = make_fixed_size_list_schema(width, /*nullable=*/false); + auto input_array = make_fixed_size_list_array(input_schema.get(), values, num_rows); + auto input = as_host_device_array(input_array); + + auto result = cudf::from_arrow_host(input_schema.get(), &input); + auto result_lists = cudf::lists_column_view{result->get_column(0)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_offsets, result_lists.offsets()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_child, result_lists.child()); + } +} + +TEST_F(FromArrowHostDeviceTest, FixedSizeListInvalidBounds) +{ + auto input_schema = make_fixed_size_list_schema(3, /*nullable=*/false); + auto input_array = make_fixed_size_list_array(input_schema.get(), {1, 2, 3}, 1); + + ArrowDeviceArray input; + memcpy(&input.array, input_array->children[0], sizeof(ArrowArray)); + input.device_id = -1; + input.device_type = ARROW_DEVICE_CPU; + input.sync_event = nullptr; + + input.array.offset = -1; + EXPECT_THROW(cudf::from_arrow_host_column(input_schema->children[0], &input), + std::invalid_argument); + + input.array.offset = 0; + input.array.length = -1; + EXPECT_THROW(cudf::from_arrow_host_column(input_schema->children[0], &input), + std::invalid_argument); + + input.array.length = std::numeric_limits::max(); + EXPECT_THROW(cudf::from_arrow_host_column(input_schema->children[0], &input), + std::overflow_error); + + input.array.offset = std::numeric_limits::max(); + input.array.length = 1; + EXPECT_THROW(cudf::from_arrow_host_column(input_schema->children[0], &input), + std::overflow_error); + + input.array.offset = std::numeric_limits::max() / 3 + 1; + input.array.length = 1; + EXPECT_THROW(cudf::from_arrow_host_column(input_schema->children[0], &input), + std::overflow_error); + + input.array.offset = 0; + input.array.length = 1; + input.array.children[0]->length = 2; + EXPECT_THROW(cudf::from_arrow_host_column(input_schema->children[0], &input), + std::invalid_argument); +} + TEST_F(FromArrowHostDeviceTest, StructColumn) { // Create cudf table diff --git a/cpp/tests/interop/from_arrow_stream_test.cpp b/cpp/tests/interop/from_arrow_stream_test.cpp index c9983f758f2e..fe06f628880f 100644 --- a/cpp/tests/interop/from_arrow_stream_test.cpp +++ b/cpp/tests/interop/from_arrow_stream_test.cpp @@ -10,11 +10,16 @@ #include #include +#include #include +#include #include #include #include +#include +#include + struct FromArrowStreamTest : public cudf::test::BaseFixture {}; void makeStreamFromArrays(std::vector arrays, @@ -125,7 +130,7 @@ namespace { // Builds a struct schema with one fixed_size_list[width] child. ArrowSchemaInitFromType // does not support NANOARROW_TYPE_FIXED_SIZE_LIST, so ArrowSchemaSetTypeFixedSize is used and // the allocated "item" child gets its type set explicitly. -nanoarrow::UniqueSchema make_fixed_size_list_stream_schema(int32_t width) +nanoarrow::UniqueSchema make_fixed_size_list_stream_schema(int32_t width, bool nullable = false) { nanoarrow::UniqueSchema schema; ArrowSchemaInit(schema.get()); @@ -134,7 +139,7 @@ nanoarrow::UniqueSchema make_fixed_size_list_stream_schema(int32_t width) NANOARROW_THROW_NOT_OK( ArrowSchemaSetTypeFixedSize(schema->children[0], NANOARROW_TYPE_FIXED_SIZE_LIST, width)); NANOARROW_THROW_NOT_OK(ArrowSchemaSetName(schema->children[0], "a")); - schema->children[0]->flags = 0; + schema->children[0]->flags = nullable ? ARROW_FLAG_NULLABLE : 0; NANOARROW_THROW_NOT_OK( ArrowSchemaSetType(schema->children[0]->children[0], NANOARROW_TYPE_INT64)); @@ -146,7 +151,8 @@ nanoarrow::UniqueSchema make_fixed_size_list_stream_schema(int32_t width) nanoarrow::UniqueArray make_fixed_size_list_chunk(ArrowSchema* schema, std::vector const& values, - int64_t num_rows) + int64_t num_rows, + std::vector const& validity = {}) { nanoarrow::UniqueArray array; NANOARROW_THROW_NOT_OK(ArrowArrayInitFromSchema(array.get(), schema, nullptr)); @@ -156,6 +162,17 @@ nanoarrow::UniqueArray make_fixed_size_list_chunk(ArrowSchema* schema, auto* list_array = array->children[0]; list_array->length = num_rows; list_array->null_count = 0; + if (!validity.empty()) { + ArrowBitmap bitmap; + ArrowBitmapInit(&bitmap); + NANOARROW_THROW_NOT_OK(ArrowBitmapReserve(&bitmap, validity.size())); + ArrowBitmapAppendInt8Unsafe( + &bitmap, reinterpret_cast(validity.data()), validity.size()); + ArrowArraySetValidityBitmap(list_array, &bitmap); + list_array->null_count = + num_rows - + ArrowBitCountSet(ArrowArrayValidityBitmap(list_array)->buffer.data, 0, validity.size()); + } auto* values_array = list_array->children[0]; NANOARROW_THROW_NOT_OK(ArrowBufferAppend(ArrowArrayBuffer(values_array, 1), @@ -207,3 +224,85 @@ TEST_F(FromArrowStreamTest, FixedSizeListChunkedTest) auto result = cudf::from_arrow_stream(&stream); CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table_view, result->view()); } + +TEST_F(FromArrowStreamTest, FixedSizeListChunkedNullsTest) +{ + constexpr cudf::size_type num_rows = 4; + auto schema = make_fixed_size_list_stream_schema(2, /*nullable=*/true); + + std::vector arrays; + arrays.push_back(make_fixed_size_list_chunk(schema.get(), {1, 2, 3, 4}, 2, /*validity=*/{1, 0})); + arrays.push_back(make_fixed_size_list_chunk(schema.get(), {5, 6, 7, 8}, 2, /*validity=*/{0, 1})); + + auto child = cudf::test::fixed_width_column_wrapper{1, 2, 3, 4, 5, 6, 7, 8}.release(); + auto offsets = cudf::test::fixed_width_column_wrapper{0, 2, 4, 6, 8}.release(); + std::vector validity{1, 0, 0, 1}; + auto [null_mask, null_count] = + cudf::test::detail::make_null_mask(validity.begin(), validity.end()); + auto expected = cudf::make_lists_column( + num_rows, std::move(offsets), std::move(child), null_count, std::move(null_mask)); + + ArrowArrayStream stream; + makeStreamFromArrays(std::move(arrays), std::move(schema), &stream); + + auto result = cudf::from_arrow_stream(&stream); + auto result_lists = cudf::lists_column_view{result->get_column(0)}; + EXPECT_EQ(result_lists.null_count(), 2); + + auto expected_logical = cudf::purge_nonempty_nulls(expected->view()); + auto result_logical = cudf::purge_nonempty_nulls(result_lists.parent()); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(*expected_logical, *result_logical); +} + +TEST_F(FromArrowStreamTest, FixedSizeListSlicedTest) +{ + constexpr cudf::size_type width = 2; + auto schema = make_fixed_size_list_stream_schema(width); + + std::vector arrays; + arrays.push_back( + make_fixed_size_list_chunk(schema.get(), {1, 2, 3, 4, 5, 6, 7, 8}, /*num_rows=*/4)); + arrays.front()->length = 2; + arrays.front()->children[0]->offset = 1; + arrays.front()->children[0]->length = 2; + + auto expected = cudf::test::lists_column_wrapper{{3, 4}, {5, 6}}; + + ArrowArrayStream stream; + makeStreamFromArrays(std::move(arrays), std::move(schema), &stream); + + auto result = cudf::from_arrow_stream(&stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->get_column(0)); +} + +TEST_F(FromArrowStreamTest, FixedSizeListBoundaryAndMultiBlockTest) +{ + constexpr cudf::size_type width = 2; + auto schema = make_fixed_size_list_stream_schema(width); + + std::vector arrays; + std::vector expected_values; + std::vector expected_offsets{0}; + for (auto const num_rows : {cudf::size_type{1024}, cudf::size_type{1025}}) { + std::vector values(num_rows * width); + std::iota(values.begin(), values.end(), static_cast(expected_values.size())); + expected_values.insert(expected_values.end(), values.begin(), values.end()); + arrays.push_back(make_fixed_size_list_chunk(schema.get(), values, num_rows)); + for (cudf::size_type i = 0; i < num_rows; ++i) { + expected_offsets.push_back(expected_offsets.back() + width); + } + } + + auto expected_offsets_col = cudf::test::fixed_width_column_wrapper( + expected_offsets.begin(), expected_offsets.end()); + auto expected_child = + cudf::test::fixed_width_column_wrapper(expected_values.begin(), expected_values.end()); + + ArrowArrayStream stream; + makeStreamFromArrays(std::move(arrays), std::move(schema), &stream); + + auto result = cudf::from_arrow_stream(&stream); + auto result_lists = cudf::lists_column_view{result->get_column(0)}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_offsets_col, result_lists.offsets()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_child, result_lists.child()); +} diff --git a/python/pylibcudf/tests/test_table.py b/python/pylibcudf/tests/test_table.py index 65c1c47186a6..18b73c090348 100644 --- a/python/pylibcudf/tests/test_table.py +++ b/python/pylibcudf/tests/test_table.py @@ -107,11 +107,13 @@ def test_from_arrow_zero_column_preserves_num_rows(): assert tbl.num_rows() == 5 -def test_from_arrow_fixed_size_list_normalizes_to_list(): - fixed = pa.array( - [[1, 2, 3], [4, 5, 6]], type=pa.list_(pa.int64(), list_size=3) - ) - expected = pa.array([[1, 2, 3], [4, 5, 6]], type=pa.list_(pa.int64())) +@pytest.mark.parametrize( + "values", + [[], [None, None], [[1, 2, 3]], [[1, 2, 3], [4, 5, 6]]], +) +def test_from_arrow_fixed_size_list_normalizes_to_list(values): + fixed = pa.array(values, type=pa.list_(pa.int64(), list_size=3)) + expected = pa.array(values, type=pa.list_(pa.int64())) column = plc.Column.from_arrow(fixed) assert column.to_arrow().equals(expected) @@ -120,6 +122,22 @@ def test_from_arrow_fixed_size_list_normalizes_to_list(): assert_table_eq(table.to_arrow(), pa.table({"a": expected})) +def test_from_arrow_fixed_size_list_in_mixed_table(): + fixed = pa.array( + [[1, 2, 3], [4, 5, 6]], type=pa.list_(pa.int64(), list_size=3) + ) + arrow_table = pa.table( + {"fixed": fixed, "integer": pa.array([7, 8]), "string": ["a", "b"]} + ) + expected = arrow_table.set_column( + 0, + "fixed", + pa.array([[1, 2, 3], [4, 5, 6]], type=pa.list_(pa.int64())), + ) + + assert_table_eq(plc.Table.from_arrow(arrow_table).to_arrow(), expected) + + def test_to_arrow_zero_column_preserves_num_rows(): arrow = plc.Table([], num_rows=5).to_arrow() assert arrow.num_columns == 0