wjones127 commented on code in PR #37292:
URL: https://github.com/apache/arrow/pull/37292#discussion_r1301835676
##########
cpp/src/arrow/compute/kernels/scalar_cast_nested.cc:
##########
@@ -145,6 +153,135 @@ void AddListCast(CastFunction* func) {
DCHECK_OK(func->AddKernel(SrcType::type_id, std::move(kernel)));
}
+template <typename DestType>
+struct CastFixedToVarList {
+ using dest_offset_type = typename DestType::offset_type;
+
+ 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();
+ ARROW_ASSIGN_OR_RAISE(out_array->buffers[0],
+ GetNullBitmapBuffer(in_array, ctx->memory_pool()));
+
+ const auto& in_type = checked_cast<const
FixedSizeListType&>(*in_array.type);
+ const int32_t list_size = in_type.list_size();
+
+ // Allocate a new offsets buffer
+ ARROW_ASSIGN_OR_RAISE(out_array->buffers[1],
+ ctx->Allocate(sizeof(dest_offset_type) *
(batch.length + 1)));
+ auto* offsets = out_array->GetMutableValues<dest_offset_type>(1);
+ dest_offset_type offset = 0;
+ for (int64_t i = 0; i <= batch.length; ++i) {
+ offsets[i] = offset;
+ offset += list_size;
+ }
+
+ // Handle values
+ std::shared_ptr<ArrayData> values = in_array.child_data[0].ToArrayData();
+ if (in_array.offset > 0 || (in_array.length * list_size < batch.length)) {
+ values = values->Slice(in_array.offset * list_size, in_array.length *
list_size);
+ }
+ 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>
+struct CastVarToFixedList {
+ using src_offset_type = typename SrcType::offset_type;
+
+ static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult*
out) {
+ const CastOptions& options = CastState::Get(ctx);
+
+ auto child_type = checked_cast<const
FixedSizeListType&>(*out->type()).value_type();
+
+ const ArraySpan& in_array = batch[0].array;
+
+ const auto& out_type = checked_cast<const
FixedSizeListType&>(*out->type());
+ const int32_t list_size = out_type.list_size();
+
+ // Validate lengths by comparing to the expected offsets.
+ const auto* offsets = in_array.GetValues<src_offset_type>(1);
+ src_offset_type expected_offset = offsets[0];
+ if (in_array.GetNullCount() > 0) {
+ for (int64_t i = 0; i <= batch.length; ++i) {
+ if (in_array.IsNull(i)) {
+ expected_offset += offsets[i + 1] + list_size;
+ } else {
+ if (offsets[i] != expected_offset) {
+ return Status::Invalid("ListType can only be casted to
FixedSizeListType ",
+ "if the lists are all the expected size.");
+ }
+ expected_offset += list_size;
+ }
+ }
+ } else {
+ // Don't need to check null slots if there are no nulls
+ for (int64_t i = 0; i <= batch.length; ++i) {
+ if (offsets[i] != expected_offset) {
+ return Status::Invalid("ListType can only be casted to
FixedSizeListType ",
+ "if the lists are all the expected size.");
+ }
+ expected_offset += list_size;
+ }
+ }
Review Comment:
You are right the null part wasn't handled. Switched to just checking the
ends.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
