| 539 | |
| 540 | struct StructFieldFunctor { |
| 541 | static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { |
| 542 | const auto& options = OptionsWrapper<StructFieldOptions>::Get(ctx); |
| 543 | std::shared_ptr<Array> current = MakeArray(batch[0].array.ToArrayData()); |
| 544 | |
| 545 | FieldPath field_path; |
| 546 | if (options.field_ref.IsNested() || options.field_ref.IsName()) { |
| 547 | ARROW_ASSIGN_OR_RAISE(field_path, options.field_ref.FindOne(*current->type())); |
| 548 | } else { |
| 549 | DCHECK(options.field_ref.IsFieldPath()); |
| 550 | field_path = *options.field_ref.field_path(); |
| 551 | } |
| 552 | |
| 553 | for (const auto& index : field_path.indices()) { |
| 554 | RETURN_NOT_OK(CheckIndex(index, *current->type())); |
| 555 | switch (current->type()->id()) { |
| 556 | case Type::STRUCT: { |
| 557 | const auto& struct_array = checked_cast<const StructArray&>(*current); |
| 558 | ARROW_ASSIGN_OR_RAISE( |
| 559 | current, struct_array.GetFlattenedField(index, ctx->memory_pool())); |
| 560 | break; |
| 561 | } |
| 562 | case Type::DENSE_UNION: { |
| 563 | // We implement this here instead of in DenseUnionArray since it's |
| 564 | // easiest to do via Take(), but DenseUnionArray can't rely on |
| 565 | // arrow::compute. See ARROW-8891. |
| 566 | const auto& union_array = checked_cast<const DenseUnionArray&>(*current); |
| 567 | |
| 568 | // Generate a bitmap for the offsets buffer based on the type codes buffer. |
| 569 | ARROW_ASSIGN_OR_RAISE( |
| 570 | std::shared_ptr<Buffer> take_bitmap, |
| 571 | ctx->AllocateBitmap(union_array.length() + union_array.offset())); |
| 572 | const int8_t* type_codes = union_array.raw_type_codes(); |
| 573 | const int8_t type_code = union_array.union_type()->type_codes()[index]; |
| 574 | int64_t offset = 0; |
| 575 | arrow::internal::GenerateBitsUnrolled( |
| 576 | take_bitmap->mutable_data(), union_array.offset(), union_array.length(), |
| 577 | [&] { return type_codes[offset++] == type_code; }); |
| 578 | |
| 579 | // Pass the combined buffer to Take(). |
| 580 | Datum take_indices( |
| 581 | ArrayData(int32(), union_array.length(), |
| 582 | {std::move(take_bitmap), union_array.value_offsets()}, |
| 583 | kUnknownNullCount, union_array.offset())); |
| 584 | // Do not slice the child since the indices are relative to the unsliced |
| 585 | // array. |
| 586 | ARROW_ASSIGN_OR_RAISE( |
| 587 | Datum result, |
| 588 | CallFunction("take", {union_array.field(index), std::move(take_indices)})); |
| 589 | current = result.make_array(); |
| 590 | break; |
| 591 | } |
| 592 | case Type::SPARSE_UNION: { |
| 593 | const auto& union_array = checked_cast<const SparseUnionArray&>(*current); |
| 594 | ARROW_ASSIGN_OR_RAISE(current, |
| 595 | union_array.GetFlattenedField(index, ctx->memory_pool())); |
| 596 | break; |
| 597 | } |
| 598 | default: |
nothing calls this directly
no test coverage detected