Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
199 changes: 190 additions & 9 deletions be/src/format/arrow/arrow_array_normalizer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,13 @@
#include "format/arrow/arrow_array_normalizer.h"

#include <arrow/array/array_base.h>
#include <arrow/array/array_nested.h>
#include <arrow/array/builder_base.h>
#include <arrow/array/builder_nested.h>
#include <arrow/compute/cast.h>
#include <arrow/type.h>

#include <limits>
#include <memory>

#include "common/check.h"
Expand All @@ -30,22 +34,164 @@ namespace doris {

namespace {

// The accepted counterpart of an encoding-only variant. Null when there is none.
std::shared_ptr<arrow::DataType> target_type_for(const arrow::DataType& type) {
switch (type.id()) {
// The accepted counterpart of an encoding-only variant. Null when neither this type nor any child
// needs normalization.
std::shared_ptr<arrow::DataType> target_type_for(const std::shared_ptr<arrow::DataType>& type) {
switch (type->id()) {
case arrow::Type::LARGE_STRING:
case arrow::Type::STRING_VIEW:
return arrow::utf8();
case arrow::Type::LARGE_BINARY:
case arrow::Type::BINARY_VIEW:
return arrow::binary();
case arrow::Type::DICTIONARY: {
const auto& dictionary = static_cast<const arrow::DictionaryType&>(*type);
auto nested_target = target_type_for(dictionary.value_type());
return nested_target != nullptr ? nested_target : dictionary.value_type();
}
case arrow::Type::RUN_END_ENCODED: {
const auto& encoded = static_cast<const arrow::RunEndEncodedType&>(*type);
auto nested_target = target_type_for(encoded.value_type());
return nested_target != nullptr ? nested_target : encoded.value_type();
}
case arrow::Type::LIST:
case arrow::Type::LARGE_LIST:
case arrow::Type::FIXED_SIZE_LIST: {
const auto& list = static_cast<const arrow::BaseListType&>(*type);
auto child_target = target_type_for(list.value_type());
if (child_target == nullptr) {
return nullptr;
}
auto child = list.value_field()->WithType(std::move(child_target));
if (type->id() == arrow::Type::LIST) {
return arrow::list(std::move(child));
}
if (type->id() == arrow::Type::LARGE_LIST) {
return arrow::large_list(std::move(child));
}
const auto& fixed = static_cast<const arrow::FixedSizeListType&>(*type);
return arrow::fixed_size_list(std::move(child), fixed.list_size());
}
case arrow::Type::STRUCT: {
arrow::FieldVector fields;
fields.reserve(type->num_fields());
bool changed = false;
for (const auto& field : type->fields()) {
auto child_target = target_type_for(field->type());
changed |= child_target != nullptr;
fields.push_back(child_target != nullptr ? field->WithType(std::move(child_target))
: field);
}
return changed ? arrow::struct_(fields) : nullptr;
}
case arrow::Type::MAP: {
const auto& map = static_cast<const arrow::MapType&>(*type);
auto key_target = target_type_for(map.key_type());
auto item_target = target_type_for(map.item_type());
if (key_target == nullptr && item_target == nullptr) {
return nullptr;
}
auto key = key_target != nullptr ? map.key_field()->WithType(std::move(key_target))
: map.key_field();
auto item = item_target != nullptr ? map.item_field()->WithType(std::move(item_target))
: map.item_field();
return std::make_shared<arrow::MapType>(std::move(key), std::move(item), map.keys_sorted());
}
default:
return nullptr;
}
}

bool contains_list_view(const arrow::DataType& type) {
switch (type.id()) {
case arrow::Type::LIST_VIEW:
case arrow::Type::LARGE_LIST_VIEW:
return true;
case arrow::Type::DICTIONARY:
return static_cast<const arrow::DictionaryType&>(type).value_type();
return contains_list_view(*static_cast<const arrow::DictionaryType&>(type).value_type());
case arrow::Type::RUN_END_ENCODED:
return static_cast<const arrow::RunEndEncodedType&>(type).value_type();
return contains_list_view(*static_cast<const arrow::RunEndEncodedType&>(type).value_type());
case arrow::Type::EXTENSION:
return contains_list_view(*static_cast<const arrow::ExtensionType&>(type).storage_type());
default:
return nullptr;
for (const auto& field : type.fields()) {
if (contains_list_view(*field->type())) {
return true;
}
}
return false;
}
}

template <typename OffsetType, typename ViewArray, typename ListBuilder, typename ListArray>
arrow::Result<std::shared_ptr<arrow::Array>> canonicalize_list_view(const ViewArray& source,
arrow::MemoryPool* pool) {
// Imported C stream arrays are not guaranteed to have validated ranges; copying an invalid
// range before this check could read beyond the child array.
auto validation = source.ValidateFull();
if (!validation.ok()) {
return validation;
}

// Preflight the full expansion before any builder mutation. Nested view builders and Arrow's
// NullBuilder can otherwise overflow signed lengths before their parent reports capacity.
if (contains_list_view(*source.value_type())) {
return arrow::Status::Invalid("nested list view canonicalization is not supported");
}
int64_t logical_value_count = 0;
constexpr int64_t max_value_count = std::numeric_limits<OffsetType>::max();
for (int64_t i = 0; i < source.length(); ++i) {
if (source.IsNull(i)) {
continue;
}
const int64_t value_length = source.value_length(i);
if (value_length > max_value_count - logical_value_count) {
return arrow::Status::CapacityError("list view logical values exceed output capacity");
}
logical_value_count += value_length;
}

auto child_builder_result = arrow::MakeBuilder(source.value_type(), pool);
if (!child_builder_result.ok()) {
return child_builder_result.status();
}
std::shared_ptr<arrow::ArrayBuilder> child_builder(child_builder_result.MoveValueUnsafe());
ListBuilder builder(pool, child_builder);
auto reserve_status = builder.Reserve(source.length());

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Reserve the expanded child before the append loop

Only the parent slots are reserved, so the child grows geometrically while shared ListView ranges are expanded. For example, 1,025 rows viewing the same 100,000 Int32 values produce a roughly 391 MiB value buffer, but Arrow's final doubling reallocates it from about 391 MiB to 781 MiB; Doris charges the new allocation before releasing the old one, transiently accounting roughly 1.17 GiB for that buffer alone. An exact reserve keeps the canonical and Doris value buffers near 782 MiB combined, leaving roughly 391 MiB less required headroom (other common query allocations aside). After validation, please compute a checked non-null logical total and reserve it on child_builder before this loop.

if (!reserve_status.ok()) {
return reserve_status;
}
reserve_status = child_builder->Reserve(logical_value_count);
if (!reserve_status.ok()) {
return reserve_status;
}

arrow::ArraySpan values(*source.values()->data());
for (int64_t i = 0; i < source.length(); ++i) {
if (source.IsNull(i)) {
auto append_status = builder.AppendNull();
if (!append_status.ok()) {
return append_status;
}
continue;
}
auto append_status = builder.Append();
if (!append_status.ok()) {
return append_status;
}
append_status = child_builder->AppendArraySlice(values, source.value_offset(i),

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Check aggregate lengths before copying into builders

A valid LargeListView<Null> can reach signed overflow here without allocating a huge buffer: use a Null child of length INT64_MAX and two shared ranges of length INT64_MAX - 1. ValidateFull() accepts both ranges independently, but Arrow 24's NullBuilder::AppendArraySlice performs unchecked length_ += length, so the second append overflows before LargeListBuilder checks capacity. Checking only this source's top-level total is not enough either: an outer large_list_view<large_list_view<null>> can have total 2 while the nested builder overflows on the same ranges during this call. Please compute checked non-null logical totals before any append at every copied view level, or reject nested view children before building, and add a no-allocation regression test.

source.value_length(i));
if (!append_status.ok()) {
return append_status;
}
}

std::shared_ptr<ListArray> out;
auto finish_status = builder.Finish(&out);
if (!finish_status.ok()) {
return finish_status;
}
return std::static_pointer_cast<arrow::Array>(out);
}

} // namespace
Expand All @@ -59,6 +205,8 @@ bool is_serde_acceptable_arrow_type(const arrow::DataType& type) {
case arrow::Type::BINARY_VIEW:
case arrow::Type::DICTIONARY:
case arrow::Type::RUN_END_ENCODED:
case arrow::Type::LIST_VIEW:
case arrow::Type::LARGE_LIST_VIEW:
return false;
// No Doris column can hold these, so they must not reach a serde either.
case arrow::Type::INTERVAL_MONTHS:
Expand All @@ -69,13 +217,19 @@ bool is_serde_acceptable_arrow_type(const arrow::DataType& type) {
case arrow::Type::DENSE_UNION:
return false;
default:
for (const auto& field : type.fields()) {
if (!is_serde_acceptable_arrow_type(*field->type())) {
return false;
}
}
return true;
}
}

Status normalize_arrow_array(const std::shared_ptr<arrow::Array>& arr,
Status normalize_arrow_array(const std::shared_ptr<arrow::Array>& arr, arrow::MemoryPool* pool,
std::shared_ptr<arrow::Array>* out) {
DORIS_CHECK(arr != nullptr);
DORIS_CHECK(pool != nullptr);
DORIS_CHECK(out != nullptr);

std::shared_ptr<arrow::Array> current = arr;
Expand All @@ -89,14 +243,41 @@ Status normalize_arrow_array(const std::shared_ptr<arrow::Array>& arr,
return Status::OK();
}

auto target = target_type_for(type);
// List views may share or reorder value ranges, so rebuild canonical offsets instead of
// exposing their buffers to a serde that requires contiguous list values.
if (type.id() == arrow::Type::LIST_VIEW) {
auto converted = canonicalize_list_view<int32_t, arrow::ListViewArray,
arrow::ListBuilder, arrow::ListArray>(
static_cast<const arrow::ListViewArray&>(*current), pool);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Normalize encoding wrappers below the rebuilt list

For list_view<dictionary<int8,int8>>, MakeBuilder creates a dictionary child builder. Its AppendArraySlice inserts decoded values, but Finish still emits a DictionaryArray backed by indices. The next normalizer pass sees only the outer LIST as acceptable and returns it unchanged; ADBC maps the dictionary's logical value type to TINYINT, so Array/Nullable/Number SerDes copy the one-byte index buffer. Logical values [42,43] are therefore returned as [0,1], and default Arrow validation cannot distinguish the equally wide buffers. Before this change the outer ListView was rejected. Please recursively normalize or reject encoding-only descendants before accepting a complex outer type, with ListView/LargeListView dictionary-child tests.

if (!converted.ok()) {
return Status::InternalError("ADBC: failed to normalize arrow type '{}': {}",
type.ToString(), converted.status().ToString());
}
current = converted.MoveValueUnsafe();
continue;
}
if (type.id() == arrow::Type::LARGE_LIST_VIEW) {
auto converted = canonicalize_list_view<int64_t, arrow::LargeListViewArray,
arrow::LargeListBuilder, arrow::LargeListArray>(
static_cast<const arrow::LargeListViewArray&>(*current), pool);
if (!converted.ok()) {
return Status::InternalError("ADBC: failed to normalize arrow type '{}': {}",
type.ToString(), converted.status().ToString());
}
current = converted.MoveValueUnsafe();
continue;
}

auto target = target_type_for(current->type());
if (target == nullptr) {
return Status::NotSupported(
"ADBC: arrow type '{}' cannot be materialized into a Doris column",
type.ToString());
}

auto casted = arrow::compute::Cast(*current, target);
arrow::compute::ExecContext exec_context(pool);
auto casted = arrow::compute::Cast(*current, target, arrow::compute::CastOptions::Safe(),
&exec_context);
if (!casted.ok()) {
return Status::InternalError("ADBC: failed to normalize arrow type '{}' to '{}': {}",
type.ToString(), target->ToString(),
Expand Down
12 changes: 7 additions & 5 deletions be/src/format/arrow/arrow_array_normalizer.h
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,17 @@ namespace doris {
/// emits string_view, Go-based drivers may emit large_* and dictionary. This normalizes them into
/// a shape the serdes accept.
///
/// Only top-level types are normalized. A nested type whose child is an unaccepted variant (say
/// list<large_utf8>) passes through and fails inside the serde, loudly rather than silently.
/// Encoding-only descendants are normalized recursively so a complex serde never mistakes encoded
/// child buffers for logical values. Nested list views remain unsupported because their shared
/// ranges cannot be canonicalized safely by an outer Arrow cast.

/// Whether the serdes take this Arrow type as-is.
bool is_serde_acceptable_arrow_type(const arrow::DataType& type);

/// Normalizes `arr` into a serde-acceptable shape. Returns it unchanged (no copy) when it already
/// is one, and fails with the offending type named when no accepted shape exists.
Status normalize_arrow_array(const std::shared_ptr<arrow::Array>& arr,
/// Normalizes `arr` into a serde-acceptable shape using `pool` for every conversion. Returns it
/// unchanged (no copy) when it already is one, and fails with the offending type named when no
/// accepted shape exists.
Status normalize_arrow_array(const std::shared_ptr<arrow::Array>& arr, arrow::MemoryPool* pool,
std::shared_ptr<arrow::Array>* out);

} // namespace doris
31 changes: 27 additions & 4 deletions be/src/format_v2/column_mapper.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1916,6 +1916,32 @@ static void attach_timestamp_semantics(const ColumnMapping& mapping, LocalColumn
attach_timestamp_semantics(child_mapping, &*child_it);
}
}
static Status apply_projected_file_definition_to_mapping(const ColumnDefinition& projected_field,
ColumnMapping* mapping) {
DORIS_CHECK(mapping != nullptr);
mapping->file_type = projected_field.type;
mapping->projected_file_children = projected_field.children;
for (auto& child_mapping : mapping->child_mappings) {
if (!child_mapping.file_local_id.has_value()) {
continue;
}
const auto child_it =
std::ranges::find_if(projected_field.children, [&](const ColumnDefinition& child) {
return child.file_local_id() == *child_mapping.file_local_id;
});
if (child_it == projected_field.children.end()) {
return Status::InternalError(
"Projected file type for '{}' is missing mapped child id {}",
mapping->file_column_name, *child_mapping.file_local_id);
}
// A full root projection changes every descendant's runtime shape too. Keep the recursive
// mapping in sync so a formerly pruned child cannot be mistaken for a trivial direct column.
RETURN_IF_ERROR(apply_projected_file_definition_to_mapping(*child_it, &child_mapping));
}
mapping->is_trivial = mapping_can_use_file_column_directly(*mapping);
return Status::OK();
}

// Update the mapping's file type according to the projection, and determine whether the projection
// is trivial (i.e. the projected file type is the same as the table type, so no need to
// rematerialize the complex value back to table layout after reading from file).
Expand All @@ -1934,10 +1960,7 @@ static Status apply_projection_to_mapping_file_type(const LocalColumnIndex& proj
field.children = mapping->original_file_children;
ColumnDefinition projected_field;
RETURN_IF_ERROR(project_column_definition(field, projection, &projected_field));
mapping->file_type = std::move(projected_field.type);
mapping->projected_file_children = std::move(projected_field.children);
mapping->is_trivial = mapping_can_use_file_column_directly(*mapping);
return Status::OK();
return apply_projected_file_definition_to_mapping(projected_field, mapping);
}

static const ColumnDefinition* find_file_child_by_name(
Expand Down
Loading
Loading