Contains logic for writing a single leaf node to parquet. This tracks the path from root to leaf. |writer| will be called after all of the definition/repetition values have been calculated for root_range with the calculated values. It is intended to abstract the complexity of writing the levels and values to parquet.
| 586 | /// values. It is intended to abstract the complexity of writing |
| 587 | /// the levels and values to parquet. |
| 588 | Status WritePath(ElementRange root_range, PathInfo* path_info, |
| 589 | ArrowWriteContext* arrow_context, |
| 590 | MultipathLevelBuilder::CallbackFunction writer) { |
| 591 | std::vector<ElementRange> stack(path_info->path.size()); |
| 592 | MultipathLevelBuilderResult builder_result; |
| 593 | builder_result.leaf_array = path_info->primitive_array; |
| 594 | builder_result.leaf_is_nullable = path_info->leaf_is_nullable; |
| 595 | |
| 596 | if (path_info->max_def_level == 0) { |
| 597 | // This case only occurs when there are no nullable or repeated |
| 598 | // columns in the path from the root to leaf. |
| 599 | int64_t leaf_length = builder_result.leaf_array->length(); |
| 600 | builder_result.def_rep_level_count = leaf_length; |
| 601 | builder_result.post_list_visited_elements.push_back({0, leaf_length}); |
| 602 | return writer(builder_result); |
| 603 | } |
| 604 | stack[0] = root_range; |
| 605 | RETURN_NOT_OK( |
| 606 | arrow_context->def_levels_buffer->Resize(/*new_size=*/0, /*shrink_to_fit*/ false)); |
| 607 | PathWriteContext context(arrow_context->memory_pool, arrow_context->def_levels_buffer); |
| 608 | // We should need at least this many entries so reserve the space ahead of time. |
| 609 | RETURN_NOT_OK(context.def_levels.Reserve(root_range.Size())); |
| 610 | if (path_info->max_rep_level > 0) { |
| 611 | RETURN_NOT_OK(context.rep_levels.Reserve(root_range.Size())); |
| 612 | } |
| 613 | |
| 614 | auto stack_base = &stack[0]; |
| 615 | auto stack_position = stack_base; |
| 616 | // This is the main loop for calculated rep/def levels. The nodes |
| 617 | // in the path implement a chain-of-responsibility like pattern |
| 618 | // where each node can add some number of repetition/definition |
| 619 | // levels to PathWriteContext and also delegate to the next node |
| 620 | // in the path to add values. The values are added through each Run(...) |
| 621 | // call and the choice to delegate to the next node (or return to the |
| 622 | // previous node) is communicated by the return value of Run(...). |
| 623 | // The loop terminates after the first node indicates all values in |
| 624 | // |root_range| are processed. |
| 625 | while (stack_position >= stack_base) { |
| 626 | PathInfo::Node& node = path_info->path[stack_position - stack_base]; |
| 627 | WritePathVisitor visitor = {.stack_position = stack_position, .context = &context}; |
| 628 | IterationResult result = std::visit(visitor, node); |
| 629 | |
| 630 | if (ARROW_PREDICT_FALSE(result == kError)) { |
| 631 | DCHECK(!context.last_status.ok()); |
| 632 | return context.last_status; |
| 633 | } |
| 634 | stack_position += static_cast<int>(result); |
| 635 | } |
| 636 | RETURN_NOT_OK(context.last_status); |
| 637 | builder_result.def_rep_level_count = context.def_levels.length(); |
| 638 | |
| 639 | if (context.rep_levels.length() > 0) { |
| 640 | // This case only occurs when there was a repeated element that needs to be |
| 641 | // processed. |
| 642 | builder_result.rep_levels = context.rep_levels.data(); |
| 643 | std::swap(builder_result.post_list_visited_elements, context.visited_elements); |
| 644 | // If it is possible when processing lists that all lists where empty. In this |
| 645 | // case no elements would have been added to post_list_visited_elements. By |