-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Add Arrow fixed-size-list ingress #23583
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
d7f79b7
f600d00
4ffdc93
7a65d70
921ac0c
463f61e
92f6898
1533e4d
be1578e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| */ | ||
|
|
||
|
|
@@ -10,6 +10,8 @@ | |
|
|
||
| #include <nanoarrow/nanoarrow.h> | ||
|
|
||
| #include <cstdint> | ||
|
|
||
| namespace cudf { | ||
| namespace detail { | ||
|
|
||
|
|
@@ -28,6 +30,53 @@ 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 Validated physical bounds for an Arrow fixed-size-list array | ||
| */ | ||
| struct fixed_size_list_layout { | ||
| int32_t width; ///< Child elements per row and LIST offset increment | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. question:
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
| 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 the INT32 LIST offset range | ||
| * | ||
| * @param arrow_view SchemaView to pull the fixed size from | ||
| * @return Number of child elements per row | ||
| */ | ||
| int32_t 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`, output row counts exceed | ||
| * `size_type`, or synthesized LIST offsets exceed INT32 | ||
| * | ||
| * @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 | ||
| * | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -30,6 +30,8 @@ | |
| #include <rmm/device_buffer.hpp> | ||
| #include <rmm/exec_policy.hpp> | ||
|
|
||
| #include <thrust/sequence.h> | ||
|
|
||
| #include <nanoarrow/nanoarrow.h> | ||
| #include <nanoarrow/nanoarrow.hpp> | ||
| #include <nanoarrow/nanoarrow_device.h> | ||
|
|
@@ -284,24 +286,60 @@ std::unique_ptr<column> dispatch_copy_from_arrow_host::operator()<cudf::struct_v | |
| input->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<std::unique_ptr<column>, 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 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), | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, I think there is a potential overflow here.
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. See prev reply. I will make this one semantically similar to prev two CUDF_EXPECTS. |
||
| layout.child_offset, | ||
| layout.child_length}; | ||
| } | ||
|
|
||
| template <> | ||
| std::unique_ptr<column> dispatch_copy_from_arrow_host::operator()<cudf::list_view>( | ||
| ArrowSchemaView const* schema, ArrowArray const* input, data_type type, bool skip_mask) | ||
| { | ||
| CUDF_EXPECTS( | ||
| input->length + 1 <= static_cast<std::int64_t>(std::numeric_limits<cudf::size_type>::max()), | ||
| "Number of rows exceeds cuDF's maximum supported row count (cudf::size_type).", | ||
| std::overflow_error); | ||
| CUDF_EXPECTS(input->length >= 0, "Number of rows must be non-negative.", std::invalid_argument); | ||
| constexpr auto max_row_count = static_cast<int64_t>(std::numeric_limits<size_type>::max()) - 1; | ||
| CUDF_EXPECTS(input->length <= max_row_count, | ||
| "Number of rows exceeds cuDF's maximum supported row count (cudf::size_type).", | ||
| std::overflow_error); | ||
|
|
||
| auto [offsets_column, offset, length] = get_offsets_column(schema, input, stream, mr); | ||
| 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); | ||
|
|
||
| ArrowSchemaView view; | ||
| NANOARROW_THROW_NOT_OK(ArrowSchemaViewInit(&view, schema->schema->children[0], nullptr)); | ||
| auto child_type = arrow_to_cudf_type(&view); | ||
|
|
||
| ArrowArray child_array(*input->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<int64_t>::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); | ||
|
|
||
|
|
@@ -394,9 +432,27 @@ std::tuple<std::unique_ptr<column>, int64_t, int64_t> copy_offsets_column( | |
|
|
||
| } // namespace | ||
|
|
||
| std::unique_ptr<column> make_fixed_size_list_offsets(size_type num_offsets, | ||
| int32_t width, | ||
| rmm::cuda_stream_view stream, | ||
| rmm::device_async_resource_ref mr) | ||
| { | ||
| auto offsets = make_numeric_column( | ||
| data_type{type_id::INT32}, num_offsets, mask_state::UNALLOCATED, stream, mr); | ||
| auto d_offsets = offsets->mutable_view().begin<int32_t>(); | ||
| thrust::sequence( | ||
| rmm::exec_policy_nosync(stream, mr), d_offsets, d_offsets + num_offsets, int32_t{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<std::unique_ptr<column>, int64_t, int64_t> get_offsets_column( | ||
| ArrowSchemaView const* schema, | ||
|
|
||


Uh oh!
There was an error while loading. Please reload this page.