[BugFix] Fix writing array column into parquet file (StarRocks#23359)

Signed-off-by: Letian Jiang <[email protected]>
Signed-off-by: Moonm3n <[email protected]>

Author: letian-jiang

be/src/formats/parquet/level_builder.cpp

@@ -113,7 +113,7 @@ void LevelBuilder::_write_column_chunk(const LevelBuilderContext& ctx, const Typ
         break;
     }
     case TYPE_ARRAY: {
-        _write_array_column_chunk_branchless(ctx, type_desc, node, col, write_leaf_callback);
+        _write_array_column_chunk(ctx, type_desc, node, col, write_leaf_callback);
         break;
     }
     case TYPE_MAP: {
@@ -399,101 +399,6 @@ void LevelBuilder::_write_array_column_chunk(const LevelBuilderContext& ctx, con
     _write_column_chunk(derived_ctx, type_desc.children[0], inner_node, elements, write_leaf_callback);
 }
 
-void LevelBuilder::_write_array_column_chunk_branchless(const LevelBuilderContext& ctx, const TypeDescriptor& type_desc,
-                                                        const ::parquet::schema::NodePtr& node, const ColumnPtr& col,
-                                                        const CallbackFunction& write_leaf_callback) {
-    DCHECK(type_desc.type == TYPE_ARRAY);
-    auto outer_node = std::static_pointer_cast<::parquet::schema::GroupNode>(node);
-    auto mid_node = std::static_pointer_cast<::parquet::schema::GroupNode>(outer_node->field(0));
-    auto inner_node = mid_node->field(0);
-
-    auto* null_col = get_raw_null_column(col);
-    auto* array_col = down_cast<ArrayColumn*>(ColumnHelper::get_data_column(col.get()));
-    const auto& elements = array_col->elements_column();
-    const auto& offsets = array_col->offsets_column()->get_data();
-
-    size_t num_levels_upper_bound = ctx._num_levels + elements->size();
-    auto def_levels = std::make_shared<std::vector<int16_t>>(num_levels_upper_bound,
-                                                             ctx._max_def_level + node->is_optional() + 1);
-    auto rep_levels = std::make_shared<std::vector<int16_t>>(num_levels_upper_bound, ctx._max_rep_level + 1);
-
-    size_t num_levels = 0; // pointer to def/rep levels
-    int offset = 0;        // pointer to current column
-    int ctx_ptr = 0;       // pointer to levels in context
-
-    if (null_col != nullptr) {
-        while (ctx_ptr < ctx._num_levels && offset < col->size()) {
-            auto def_level = ctx._def_levels ? (*ctx._def_levels)[ctx_ptr] : 0;
-            auto rep_level = ctx._rep_levels ? (*ctx._rep_levels)[ctx_ptr] : 0;
-
-            uint8_t null_in_parent_column = def_level != ctx._max_def_level;
-            uint8_t null_in_current_column = null_col[offset] == 1;
-            uint8_t empty_array = offsets[offset + 1] == offsets[offset];
-            uint8_t is_optional = node->is_optional();
-            auto array_size = offsets[offset + 1] - offsets[offset];
-
-            // set def_level back for null in parent column
-            (*def_levels)[num_levels] += null_in_parent_column * (def_level - (*def_levels)[num_levels]);
-            // decrement 1 for empty array
-            (*def_levels)[num_levels] -= ((1 - null_in_parent_column) & (1 - null_in_current_column) & empty_array);
-            // decrement node->is_optional for null in current column
-            (*def_levels)[num_levels] -= ((1 - null_in_parent_column) & null_in_current_column & is_optional);
-            (*rep_levels)[num_levels] = rep_level;
-
-            // update pointers
-            uint8_t advance_one_step = null_in_parent_column | null_in_current_column | empty_array;
-            num_levels += advance_one_step + (1 - advance_one_step) * array_size;
-            offset += (1 - null_in_parent_column);
-            ctx_ptr++;
-        }
-    } else {
-        while (ctx_ptr < ctx._num_levels && offset < col->size()) {
-            auto def_level = ctx._def_levels ? (*ctx._def_levels)[ctx_ptr] : 0;
-            auto rep_level = ctx._rep_levels ? (*ctx._rep_levels)[ctx_ptr] : 0;
-
-            uint8_t null_in_parent_column = def_level != ctx._max_def_level;
-            uint8_t empty_array = offsets[offset + 1] == offsets[offset];
-            auto array_size = offsets[offset + 1] - offsets[offset];
-
-            // set def_level back for null in parent column
-            (*def_levels)[num_levels] =
-                    null_in_parent_column * def_level + (1 - null_in_parent_column) * (*def_levels)[num_levels];
-            // decrement 1 for empty array
-            (*def_levels)[num_levels] -= ((1 - null_in_parent_column) & empty_array);
-            (*rep_levels)[num_levels] = rep_level;
-
-            // update pointers
-            uint8_t advance_one_step = null_in_parent_column | empty_array;
-            num_levels += advance_one_step + (1 - advance_one_step) * array_size;
-            offset += (1 - null_in_parent_column);
-            ctx_ptr++;
-        }
-    }
-    DCHECK(col->size() == offset);
-
-    // handle remaining undefined entries
-    if (ctx_ptr < ctx._num_levels) {
-        if (ctx._def_levels == nullptr) {
-            memset(def_levels->data() + num_levels, 0, ctx._num_levels - ctx_ptr);
-        } else {
-            memcpy(def_levels->data() + num_levels, ctx._def_levels->data() + ctx_ptr, ctx._num_levels - ctx_ptr);
-        }
-
-        if (ctx._rep_levels == nullptr) {
-            memset(rep_levels->data() + num_levels, 0, ctx._num_levels - ctx_ptr);
-        } else {
-            memcpy(rep_levels->data() + num_levels, ctx._rep_levels->data() + ctx_ptr, ctx._num_levels - ctx_ptr);
-        }
-    }
-
-    def_levels->resize(num_levels);
-    rep_levels->resize(num_levels);
-    LevelBuilderContext derived_ctx(def_levels->size(), def_levels, ctx._max_def_level + node->is_optional() + 1,
-                                    rep_levels, ctx._max_rep_level + 1);
-
-    _write_column_chunk(derived_ctx, type_desc.children[0], inner_node, elements, write_leaf_callback);
-}
-
 void LevelBuilder::_write_map_column_chunk(const LevelBuilderContext& ctx, const TypeDescriptor& type_desc,
                                            const ::parquet::schema::NodePtr& node, const ColumnPtr& col,
                                            const CallbackFunction& write_leaf_callback) {

be/src/formats/parquet/level_builder.h

@@ -126,10 +126,6 @@ class LevelBuilder {
                                    const ::parquet::schema::NodePtr& node, const ColumnPtr& col,
                                    const CallbackFunction& write_leaf_callback);
 
-    void _write_array_column_chunk_branchless(const LevelBuilderContext& ctx, const TypeDescriptor& type_desc,
-                                              const ::parquet::schema::NodePtr& node, const ColumnPtr& col,
-                                              const CallbackFunction& write_leaf_callback);
-
     void _write_map_column_chunk(const LevelBuilderContext& ctx, const TypeDescriptor& type_desc,
                                  const ::parquet::schema::NodePtr& node, const ColumnPtr& col,
                                  const CallbackFunction& write_leaf_callback);