be/src/format_v2/schema_projection.cpp

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include "format_v2/schema_projection.h"

#include <algorithm>
#include <memory>
#include <utility>

#include "core/assert_cast.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_map.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_struct.h"

namespace doris::format {
namespace {

// Rebuild the complex DataType for one already-pruned semantic ColumnDefinition node.
//
// The caller has already matched the projection against ColumnDefinition::children and preserved
// the file-local child order. This helper only mirrors those projected semantic children into the
// node type. It intentionally does not understand physical format wrappers. In particular, a MAP
// node is expected to have semantic children [key, value], even if the underlying format stores a
// wrapper such as Parquet key_value/entry.
Status rebuild_semantic_projected_type(const DataTypePtr& original_type,
                                       const std::vector<ColumnDefinition>& projected_children,
                                       DataTypePtr* projected_type) {
    DORIS_CHECK(original_type != nullptr);
    DORIS_CHECK(projected_type != nullptr);

    DataTypePtr nested_projected_type;
    const auto primitive_type = remove_nullable(original_type)->get_primitive_type();
    switch (primitive_type) {
    case TYPE_STRUCT: {
        DataTypes child_types;
        Strings child_names;
        child_types.reserve(projected_children.size());
        child_names.reserve(projected_children.size());
        for (const auto& child : projected_children) {
            child_types.push_back(child.type);
            child_names.push_back(child.name);
        }
        nested_projected_type = std::make_shared<DataTypeStruct>(child_types, child_names);
        break;
    }
    case TYPE_ARRAY:
        DORIS_CHECK(projected_children.size() == 1);
        nested_projected_type = std::make_shared<DataTypeArray>(projected_children[0].type);
        break;
    case TYPE_MAP: {
        DORIS_CHECK(remove_nullable(original_type)->get_primitive_type() == TYPE_MAP);
        const auto* original_map_type =
                assert_cast<const DataTypeMap*>(remove_nullable(original_type).get());
        DataTypePtr key_type = original_map_type->get_key_type();
        DataTypePtr value_type;
        for (const auto& child : projected_children) {
            // Partial MAP projection only prunes the value subtree. The key stream must remain
            // complete because it defines entry existence and offsets when materializing ColumnMap;
            // the projected DataTypeMap also preserves the original key type instead of rebuilding
            // it from children. If a caller includes key in the semantic child list, ignore it
            // here; the presence of a value child still decides the projected value shape.
            if (child.file_local_id() == 0 || child.name == "key") {
                continue;
            }
            if (child.file_local_id() == 1 || child.name == "value") {
                value_type = child.type;
            }
        }
        if (value_type == nullptr) {
            return Status::NotSupported("MAP projection for type {} contains no value child",
                                        original_type->get_name());
        }
        nested_projected_type = std::make_shared<DataTypeMap>(key_type, value_type);
        break;
    }
    default:
        return Status::InvalidArgument("Cannot project children from non-complex type {}",
                                       original_type->get_name());
    }

    *projected_type = original_type->is_nullable() ? make_nullable(nested_projected_type)
                                                   : nested_projected_type;
    return Status::OK();
}

} // namespace

Status project_column_definition(const ColumnDefinition& field, const LocalColumnIndex& projection,
                                 ColumnDefinition* projected_field) {
    if (projected_field == nullptr) {
        return Status::InvalidArgument("projected_field is null");
    }
    *projected_field = field;
    if (projection.project_all_children || projection.children.empty()) {
        return Status::OK();
    }

    projected_field->children.clear();
    for (const auto& child_projection : projection.children) {
        if (child_projection.local_id() == -1) {
            return Status::InvalidArgument("Empty projection path for field {}", field.name);
        }
        const auto child_it =
                std::ranges::find_if(field.children, [&](const ColumnDefinition& child) {
                    return child.file_local_id() == child_projection.local_id();
                });
        if (child_it == field.children.end()) {
            return Status::InvalidArgument("Invalid projection child id {} for field {}",
                                           child_projection.local_id(), field.name);
        }
    }
    for (const auto& child : field.children) {
        const auto child_projection_it =
                std::ranges::find_if(projection.children, [&](const LocalColumnIndex& child_proj) {
                    return child_proj.local_id() == child.file_local_id();
                });
        if (child_projection_it == projection.children.end()) {
            continue;
        }
        ColumnDefinition projected_child;
        RETURN_IF_ERROR(project_column_definition(child, *child_projection_it, &projected_child));
        projected_field->children.push_back(std::move(projected_child));
    }
    if (projected_field->children.empty()) {
        return Status::NotSupported("Projection for field {} contains no children", field.name);
    }

    return rebuild_semantic_projected_type(field.type, projected_field->children,
                                           &projected_field->type);
}

} // namespace doris::format