-
Notifications
You must be signed in to change notification settings - Fork 4.3k
GH-50994: [C++][Compute] Implement casting from ListView to List with zero-copy fast-path #50976
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
125ab1c
289afef
53085f0
beb87f7
ce16dcc
68def9d
5e4cf34
6bd6303
5d7c759
c929f7e
b312f28
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 |
|---|---|---|
|
|
@@ -26,8 +26,10 @@ | |
| #include "arrow/array/builder_primitive.h" | ||
| #include "arrow/compute/api_scalar.h" | ||
| #include "arrow/compute/cast.h" | ||
| #include "arrow/compute/exec.h" | ||
| #include "arrow/compute/kernels/common_internal.h" | ||
| #include "arrow/compute/kernels/scalar_cast_internal.h" | ||
| #include "arrow/util/bit_run_reader.h" | ||
| #include "arrow/util/bitmap_ops.h" | ||
| #include "arrow/util/int_util.h" | ||
| #include "arrow/util/logging_internal.h" | ||
|
|
@@ -141,6 +143,146 @@ void AddListCast(CastFunction* func) { | |
| DCHECK_OK(func->AddKernel(SrcType::type_id, std::move(kernel))); | ||
| } | ||
|
|
||
| // (Large)ListView<T> -> (Large)List<U> | ||
| template <typename SrcType, typename DestType> | ||
| struct CastListViewToVarList { | ||
| using src_offset_type = typename SrcType::offset_type; | ||
| using dest_offset_type = typename DestType::offset_type; | ||
|
|
||
| static constexpr bool is_downcast = sizeof(src_offset_type) > sizeof(dest_offset_type); | ||
|
|
||
| static bool IsContiguous(const ArraySpan& in_array) { | ||
| const auto* offsets = in_array.GetValues<src_offset_type>(1); | ||
| const auto* sizes = in_array.GetValues<src_offset_type>(2); | ||
| for (int64_t i = 0; i < in_array.length - 1; ++i) { | ||
|
Member
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. You could perhaps use |
||
| if (in_array.IsNull(i) && sizes[i] != 0) { | ||
| return false; | ||
| } | ||
| if (offsets[i] + sizes[i] != offsets[i + 1]) { | ||
|
Member
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. If the entry is null or zero-sized, then the exact value of (this is not a bug of course, just an additional optimization opportunity) |
||
| return false; | ||
| } | ||
| } | ||
| if (in_array.length > 0 && in_array.IsNull(in_array.length - 1) && | ||
| sizes[in_array.length - 1] != 0) { | ||
| return false; | ||
| } | ||
| return true; | ||
| } | ||
|
|
||
| static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { | ||
| const CastOptions& options = CastState::Get(ctx); | ||
| auto child_type = checked_cast<const DestType&>(*out->type()).value_type(); | ||
| const ArraySpan& in_array = batch[0].array; | ||
| ArrayData* out_array = out->array_data().get(); | ||
|
|
||
| DCHECK_NE(in_array.length, 0); | ||
|
|
||
| ARROW_ASSIGN_OR_RAISE(out_array->buffers[0], | ||
| GetOrCopyNullBitmapBuffer(in_array, ctx->memory_pool())); | ||
|
|
||
| std::shared_ptr<ArrayData> values = in_array.child_data[0].ToArrayData(); | ||
|
|
||
| const auto* offsets = in_array.GetValues<src_offset_type>(1); | ||
| const auto* sizes = in_array.GetValues<src_offset_type>(2); | ||
|
|
||
| // Allocate destination offsets buffer (shared by both paths) | ||
| ARROW_ASSIGN_OR_RAISE(out_array->buffers[1], ctx->Allocate(sizeof(dest_offset_type) * | ||
| (in_array.length + 1))); | ||
| auto* dest_offsets = out_array->GetMutableValues<dest_offset_type>(1); | ||
|
|
||
| if (IsContiguous(in_array)) { | ||
|
HuaHuaY marked this conversation as resolved.
|
||
| // Zero-copy fast-path: shift offsets and slice child values | ||
| src_offset_type start_offset = offsets[0]; | ||
| src_offset_type abs_end_offset = | ||
| offsets[in_array.length - 1] + sizes[in_array.length - 1]; | ||
|
|
||
| if constexpr (is_downcast) { | ||
| src_offset_type range = abs_end_offset - start_offset; | ||
| if (range > std::numeric_limits<dest_offset_type>::max()) { | ||
| return Status::Invalid("Array of type ", in_array.type->ToString(), | ||
| " too large to convert to ", | ||
| out_array->type->ToString()); | ||
| } | ||
| } | ||
|
|
||
| for (int64_t i = 0; i < in_array.length; ++i) { | ||
| dest_offsets[i] = static_cast<dest_offset_type>(offsets[i] - start_offset); | ||
| } | ||
| dest_offsets[in_array.length] = | ||
| static_cast<dest_offset_type>(abs_end_offset - start_offset); | ||
|
|
||
| values = values->Slice(start_offset, abs_end_offset - start_offset); | ||
| } else { | ||
| // Non-contiguous path: compute new offsets using SetBitRunReader for bitmap | ||
| // traversal | ||
| src_offset_type current_offset = 0; | ||
| dest_offsets[0] = 0; | ||
| const uint8_t* validity = in_array.buffers[0].data; | ||
|
|
||
| if (validity == nullptr) { | ||
| for (int64_t i = 0; i < in_array.length; ++i) { | ||
| current_offset += sizes[i]; | ||
| dest_offsets[i + 1] = static_cast<dest_offset_type>(current_offset); | ||
| } | ||
| } else { | ||
| arrow::internal::SetBitRunReader reader(validity, in_array.offset, | ||
| in_array.length); | ||
| int64_t last_idx = 0; | ||
| while (true) { | ||
| const auto run = reader.NextRun(); | ||
| if (run.length == 0) { | ||
| break; | ||
| } | ||
| for (int64_t i = last_idx; i < run.position; ++i) { | ||
| dest_offsets[i + 1] = static_cast<dest_offset_type>(current_offset); | ||
| } | ||
| for (int64_t i = run.position; i < run.position + run.length; ++i) { | ||
| current_offset += sizes[i]; | ||
| dest_offsets[i + 1] = static_cast<dest_offset_type>(current_offset); | ||
| } | ||
| last_idx = run.position + run.length; | ||
| } | ||
| for (int64_t i = last_idx; i < in_array.length; ++i) { | ||
| dest_offsets[i + 1] = static_cast<dest_offset_type>(current_offset); | ||
| } | ||
| } | ||
|
Comment on lines
+218
to
+248
Member
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. I think you could simplify the implementation by having the same loop offsets for both branches. That loop would compute all destination offsets and compute whether the source entries are contiguous, all in one go. The contiguity information is mostly useful to know how to compute |
||
|
|
||
| if constexpr (is_downcast) { | ||
| if (current_offset > std::numeric_limits<dest_offset_type>::max()) { | ||
| return Status::Invalid("Array of type ", in_array.type->ToString(), | ||
| " too large to convert to ", | ||
| out_array->type->ToString()); | ||
| } | ||
| } | ||
|
|
||
| // Use TypeTraits<SrcType>::ArrayType::Flatten to get values | ||
| using ArrayType = typename TypeTraits<SrcType>::ArrayType; | ||
| auto input_array = MakeArray(in_array.ToArrayData()); | ||
| const auto& list_view_array = checked_cast<const ArrayType&>(*input_array); | ||
| ARROW_ASSIGN_OR_RAISE(auto flattened, list_view_array.Flatten(ctx->memory_pool())); | ||
| values = flattened->data(); | ||
| } | ||
|
|
||
| // Cast values | ||
| ARROW_ASSIGN_OR_RAISE(Datum cast_values, | ||
| Cast(values, child_type, options, ctx->exec_context())); | ||
| DCHECK(cast_values.is_array()); | ||
| out_array->child_data.push_back(cast_values.array()); | ||
|
|
||
| return Status::OK(); | ||
| } | ||
| }; | ||
|
|
||
| template <typename SrcType, typename DestType> | ||
| void AddListViewCast(CastFunction* func) { | ||
| ScalarKernel kernel; | ||
| kernel.exec = CastListViewToVarList<SrcType, DestType>::Exec; | ||
| kernel.signature = | ||
| KernelSignature::Make({InputType(SrcType::type_id)}, kOutputTargetType); | ||
| kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE; | ||
| DCHECK_OK(func->AddKernel(SrcType::type_id, std::move(kernel))); | ||
| } | ||
|
|
||
| template <typename DestType> | ||
| struct CastFixedToVarList { | ||
| using dest_offset_type = typename DestType::offset_type; | ||
|
|
@@ -487,18 +629,18 @@ std::vector<std::shared_ptr<CastFunction>> GetNestedCasts() { | |
| auto cast_list = std::make_shared<CastFunction>("cast_list", Type::LIST); | ||
| AddCommonCasts(Type::LIST, kOutputTargetType, cast_list.get()); | ||
| AddListCast<ListType, ListType>(cast_list.get()); | ||
| AddListCast<ListViewType, ListType>(cast_list.get()); | ||
| AddListViewCast<ListViewType, ListType>(cast_list.get()); | ||
| AddListCast<LargeListType, ListType>(cast_list.get()); | ||
| AddListCast<LargeListViewType, ListType>(cast_list.get()); | ||
| AddListViewCast<LargeListViewType, ListType>(cast_list.get()); | ||
| AddTypeToTypeCast<CastFixedToVarList<ListType>, FixedSizeListType>(cast_list.get()); | ||
|
|
||
| auto cast_large_list = | ||
| std::make_shared<CastFunction>("cast_large_list", Type::LARGE_LIST); | ||
| AddCommonCasts(Type::LARGE_LIST, kOutputTargetType, cast_large_list.get()); | ||
| AddListCast<ListType, LargeListType>(cast_large_list.get()); | ||
| AddListCast<ListViewType, LargeListType>(cast_large_list.get()); | ||
| AddListViewCast<ListViewType, LargeListType>(cast_large_list.get()); | ||
| AddListCast<LargeListType, LargeListType>(cast_large_list.get()); | ||
| AddListCast<LargeListViewType, LargeListType>(cast_large_list.get()); | ||
| AddListViewCast<LargeListViewType, LargeListType>(cast_large_list.get()); | ||
| AddTypeToTypeCast<CastFixedToVarList<LargeListType>, FixedSizeListType>( | ||
| cast_large_list.get()); | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3646,6 +3646,94 @@ TEST(Cast, ListToListOptionsPassthru) { | |
| } | ||
| } | ||
|
|
||
| TEST(Cast, ListViewToList) { | ||
|
HuaHuaY marked this conversation as resolved.
|
||
| // 1. Contiguous ListView (with nulls) | ||
| auto contiguous_src = | ||
| ArrayFromJSON(list_view(int16()), "[[10, 20], null, [30], [40, 50]]"); | ||
| auto contiguous_expected = | ||
| ArrayFromJSON(list(int16()), "[[10, 20], null, [30], [40, 50]]"); | ||
| CheckCast(contiguous_src, contiguous_expected); | ||
|
|
||
| // Assert zero-copy for contiguous values | ||
| ASSERT_OK_AND_ASSIGN(auto cast_result, Cast(contiguous_src, list(int16()))); | ||
|
Member
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. Can you call |
||
| auto src_lv = std::dynamic_pointer_cast<ListViewArray>(contiguous_src); | ||
| auto res_list = std::dynamic_pointer_cast<ListArray>(cast_result.make_array()); | ||
| ASSERT_EQ(res_list->values()->data()->buffers[1]->address(), | ||
| src_lv->values()->data()->buffers[1]->address()); | ||
| ASSERT_OK(res_list->ValidateFull()); | ||
|
|
||
| // 2. Gapped/Non-contiguous ListView | ||
| auto values = ArrayFromJSON(int16(), "[10, 20, 999, 30, 40, 50]"); | ||
| auto offsets = ArrayFromJSON(int32(), "[0, 3]"); | ||
| auto sizes = ArrayFromJSON(int32(), "[2, 3]"); | ||
| ASSERT_OK_AND_ASSIGN(auto gapped_src, | ||
| ListViewArray::FromArrays(*offsets, *sizes, *values)); | ||
| auto gapped_expected = ArrayFromJSON(list(int16()), "[[10, 20], [30, 40, 50]]"); | ||
| CheckCast(gapped_src, gapped_expected); | ||
|
|
||
| // 3. Overlapping ListView | ||
| auto overlapping_offsets = ArrayFromJSON(int32(), "[0, 1]"); | ||
| auto overlapping_sizes = ArrayFromJSON(int32(), "[2, 2]"); | ||
| ASSERT_OK_AND_ASSIGN( | ||
| auto overlapping_src, | ||
| ListViewArray::FromArrays(*overlapping_offsets, *overlapping_sizes, *values)); | ||
| auto overlapping_expected = ArrayFromJSON(list(int16()), "[[10, 20], [20, 999]]"); | ||
| CheckCast(overlapping_src, overlapping_expected); | ||
|
|
||
| // 4. Large ListView to List and vice versa (with nulls) | ||
| auto large_contiguous_src = | ||
| ArrayFromJSON(large_list_view(int16()), "[[10, 20], null, [30], [40, 50]]"); | ||
| auto large_contiguous_expected = | ||
| ArrayFromJSON(large_list(int16()), "[[10, 20], null, [30], [40, 50]]"); | ||
| CheckCast(large_contiguous_src, large_contiguous_expected); | ||
| CheckCast(contiguous_src, large_contiguous_expected); | ||
| CheckCast(large_contiguous_src, contiguous_expected); | ||
|
|
||
| // 5. Null Propagation | ||
|
Member
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. Can you add nulls in the examples above? This will probably stress more situations. |
||
| auto nulls_src = ArrayFromJSON(list_view(int16()), "[[10, null], null, [40, 50]]"); | ||
| auto nulls_expected = ArrayFromJSON(list(int16()), "[[10, null], null, [40, 50]]"); | ||
| CheckCast(nulls_src, nulls_expected); | ||
|
|
||
| // 6. Generic and Nested Type casting | ||
| auto string_src = | ||
| ArrayFromJSON(list_view(utf8()), "[[\"a\", \"b\"], [\"c\"], [\"d\", \"e\"]]"); | ||
| auto string_expected = | ||
| ArrayFromJSON(list(utf8()), "[[\"a\", \"b\"], [\"c\"], [\"d\", \"e\"]]"); | ||
| CheckCast(string_src, string_expected); | ||
|
|
||
| auto type_change_src = ArrayFromJSON(list_view(int16()), "[[10, 20], [30], [40, 50]]"); | ||
| auto type_change_expected = ArrayFromJSON(list(int32()), "[[10, 20], [30], [40, 50]]"); | ||
| CheckCast(type_change_src, type_change_expected); | ||
|
|
||
| // 7. Non-Contiguous Slice Boundary Verification | ||
| auto sliced_gapped_src = gapped_src->Slice(1, 1); | ||
| auto sliced_gapped_expected = ArrayFromJSON(list(int16()), "[[30, 40, 50]]"); | ||
| CheckCast(sliced_gapped_src, sliced_gapped_expected); | ||
|
|
||
| auto sliced_overlapping_src = overlapping_src->Slice(1, 1); | ||
| auto sliced_overlapping_expected = ArrayFromJSON(list(int16()), "[[20, 999]]"); | ||
| CheckCast(sliced_overlapping_src, sliced_overlapping_expected); | ||
|
|
||
| // 8. ListView with Nulls containing overflow values in null slots | ||
| auto null_val_src_values = ArrayFromJSON(int32(), "[10, 40000]"); | ||
| auto null_val_src_offsets = ArrayFromJSON(int32(), "[0, 1]"); | ||
| auto null_val_src_sizes = ArrayFromJSON(int32(), "[1, 1]"); | ||
| ASSERT_OK_AND_ASSIGN(auto null_val_src, ListViewArray::FromArrays( | ||
| *null_val_src_offsets, *null_val_src_sizes, | ||
| *null_val_src_values)); | ||
| auto null_val_src_masked = MaskArrayWithNullsAt(null_val_src, {1}); | ||
| auto null_val_expected = ArrayFromJSON(list(int16()), "[[10], null]"); | ||
| CheckCast(null_val_src_masked, null_val_expected); | ||
|
Member
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. Can you add a test with zero-length inputs?
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. Thanks for the detailed review @pitrou! I'll address all the mandatory items:
The SetBitRunReader optimization and unified loop are noted as future improvements. Will push the fixes very shortly! |
||
|
|
||
| // 9. Zero-length ListView input | ||
| auto empty_src = ArrayFromJSON(list_view(int16()), "[]"); | ||
| auto empty_expected = ArrayFromJSON(list(int16()), "[]"); | ||
| CheckCast(empty_src, empty_expected); | ||
| auto large_empty_src = ArrayFromJSON(large_list_view(int32()), "[]"); | ||
| auto large_empty_expected = ArrayFromJSON(large_list(int32()), "[]"); | ||
| CheckCast(large_empty_src, large_empty_expected); | ||
| } | ||
|
|
||
| static void CheckFSLToFSL(const std::vector<std::shared_ptr<DataType>>& value_types, | ||
| const std::string& json_data, | ||
| const std::string& tweaked_val_bit_string, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's add a comment summarizing this:
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added // (Large)ListView -> (Large)List comment before the template