diff --git a/be/src/core/data_type_serde/data_type_array_serde.cpp b/be/src/core/data_type_serde/data_type_array_serde.cpp index bef17ab5a58227..e2ecf2cb5e0536 100644 --- a/be/src/core/data_type_serde/data_type_array_serde.cpp +++ b/be/src/core/data_type_serde/data_type_array_serde.cpp @@ -321,6 +321,77 @@ Status DataTypeArraySerDe::write_column_to_arrow(const IColumn& column, const Nu return Status::OK(); } +namespace { + +template +Status write_array_column_to_target(const IColumn& column, const NullMap* null_map, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + WriteNested&& write_nested) { + const auto& array_column = assert_cast(column); + const auto& offsets = array_column.get_offsets(); + const auto& nested_data = array_column.get_data(); + auto& builder = assert_cast(*array_builder); + auto* nested_builder = builder.value_builder(); + for (size_t array_idx = start; array_idx < end; ++array_idx) { + if (null_map != nullptr && (*null_map)[array_idx]) { + RETURN_IF_ERROR(checkArrowStatus(builder.AppendNull(), column, *array_builder)); + continue; + } + RETURN_IF_ERROR(checkArrowStatus(builder.Append(), column, *array_builder)); + RETURN_IF_ERROR(write_nested(nested_data, nested_builder, offsets[array_idx - 1], + offsets[array_idx])); + } + return Status::OK(); +} + +} // namespace + +Status DataTypeArraySerDe::write_column_to_paimon_arrow( + const std::shared_ptr& type, const IColumn& column, + const NullMap* null_map, const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const { + // Reject an incompatible target before casting its nested schema or builder. + if (field->type()->id() != arrow::Type::LIST || + array_builder->type()->id() != arrow::Type::LIST) { + return Status::InvalidArgument("Paimon array writer requires an Arrow list field"); + } + const auto& array_type = assert_cast(*type); + const auto& list_type = assert_cast(*field->type()); + const auto& nested_field = list_type.value_field(); + return write_array_column_to_target( + column, null_map, array_builder, start, end, + [&](const IColumn& nested_data, arrow::ArrayBuilder* nested_builder, + int64_t nested_start, int64_t nested_end) { + return nested_serde->write_column_to_paimon_arrow( + array_type.get_nested_type(), nested_data, nullptr, nested_field, + nested_builder, nested_start, nested_end, ctz); + }); +} + +Status DataTypeArraySerDe::write_column_to_iceberg_arrow( + const std::shared_ptr& type, const IColumn& column, + const NullMap* null_map, const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const { + // Reject an incompatible target before casting its nested schema or builder. + if (field->type()->id() != arrow::Type::LIST || + array_builder->type()->id() != arrow::Type::LIST) { + return Status::InvalidArgument("Iceberg array writer requires an Arrow list field"); + } + const auto& array_type = assert_cast(*type); + const auto& list_type = assert_cast(*field->type()); + const auto& nested_field = list_type.value_field(); + return write_array_column_to_target( + column, null_map, array_builder, start, end, + [&](const IColumn& nested_data, arrow::ArrayBuilder* nested_builder, + int64_t nested_start, int64_t nested_end) { + return nested_serde->write_column_to_iceberg_arrow( + array_type.get_nested_type(), nested_data, nullptr, nested_field, + nested_builder, nested_start, nested_end, ctz); + }); +} + Status DataTypeArraySerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { diff --git a/be/src/core/data_type_serde/data_type_array_serde.h b/be/src/core/data_type_serde/data_type_array_serde.h index 4a74649336e586..eaf3644efc6879 100644 --- a/be/src/core/data_type_serde/data_type_array_serde.h +++ b/be/src/core/data_type_serde/data_type_array_serde.h @@ -92,6 +92,16 @@ class DataTypeArraySerDe : public DataTypeSerDe { Status write_column_to_arrow(const IColumn& column, const NullMap* null_map, arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_paimon_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_iceberg_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_orc(IColumn& column, const OrcDecodedColumnView& view) const override; diff --git a/be/src/core/data_type_serde/data_type_map_serde.cpp b/be/src/core/data_type_serde/data_type_map_serde.cpp index e1d462144c5ee6..1feec24aa5841b 100644 --- a/be/src/core/data_type_serde/data_type_map_serde.cpp +++ b/be/src/core/data_type_serde/data_type_map_serde.cpp @@ -26,6 +26,7 @@ #include "core/column/column.h" #include "core/column/column_const.h" #include "core/column/column_map.h" +#include "core/data_type/data_type_map.h" #include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/complex_type_deserialize_util.h" #include "core/data_type_serde/orc_serde_utils.h" @@ -440,6 +441,108 @@ Status DataTypeMapSerDe::write_column_to_arrow(const IColumn& column, const Null return Status::OK(); } +namespace { + +template +Status write_map_column_to_target(const IColumn& column, const NullMap* null_map, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + WriteKey&& write_key, WriteValue&& write_value) { + auto& builder = assert_cast(*array_builder); + const auto& map_column = assert_cast(column); + const IColumn& nested_keys_column = map_column.get_keys(); + const IColumn& nested_values_column = map_column.get_values(); + DCHECK(nested_keys_column.is_nullable()); + DCHECK(nested_values_column.is_nullable()); + const auto* keys_nullmap_data = + check_and_get_column(nested_keys_column)->get_null_map_data().data(); + const auto& offsets = map_column.get_offsets(); + auto* key_builder = builder.key_builder(); + auto* value_builder = builder.item_builder(); + + for (size_t row = start; row < end; ++row) { + if (null_map != nullptr && (*null_map)[row]) { + RETURN_IF_ERROR(checkArrowStatus(builder.AppendNull(), column, *array_builder)); + continue; + } + if (simd::contain_one(keys_nullmap_data + offsets[row - 1], + offsets[row] - offsets[row - 1])) { + return Status::Error(ErrorCode::INVALID_ARGUMENT, + "Can not write null value of map key to arrow."); + } + RETURN_IF_ERROR(checkArrowStatus(builder.Append(), column, *array_builder)); + RETURN_IF_ERROR(write_key(nested_keys_column, key_builder, offsets[row - 1], offsets[row])); + RETURN_IF_ERROR( + write_value(nested_values_column, value_builder, offsets[row - 1], offsets[row])); + } + return Status::OK(); +} + +} // namespace + +Status DataTypeMapSerDe::write_column_to_paimon_arrow(const std::shared_ptr& type, + const IColumn& column, + const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, + int64_t start, int64_t end, + const cctz::time_zone& ctz) const { + // Reject an incompatible target before casting its nested schema or builder. + if (field->type()->id() != arrow::Type::MAP || + array_builder->type()->id() != arrow::Type::MAP) { + return Status::InvalidArgument("Paimon map writer requires an Arrow map field"); + } + const auto& map_type = assert_cast(*type); + const auto& arrow_map_type = assert_cast(*field->type()); + const auto& key_field = arrow_map_type.key_field(); + const auto& value_field = arrow_map_type.item_field(); + return write_map_column_to_target( + column, null_map, array_builder, start, end, + [&](const IColumn& nested_keys, arrow::ArrayBuilder* key_builder, int64_t nested_start, + int64_t nested_end) { + return key_serde->write_column_to_paimon_arrow(map_type.get_key_type(), nested_keys, + nullptr, key_field, key_builder, + nested_start, nested_end, ctz); + }, + [&](const IColumn& nested_values, arrow::ArrayBuilder* value_builder, + int64_t nested_start, int64_t nested_end) { + return value_serde->write_column_to_paimon_arrow( + map_type.get_value_type(), nested_values, nullptr, value_field, + value_builder, nested_start, nested_end, ctz); + }); +} + +Status DataTypeMapSerDe::write_column_to_iceberg_arrow(const std::shared_ptr& type, + const IColumn& column, + const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, + int64_t start, int64_t end, + const cctz::time_zone& ctz) const { + // Reject an incompatible target before casting its nested schema or builder. + if (field->type()->id() != arrow::Type::MAP || + array_builder->type()->id() != arrow::Type::MAP) { + return Status::InvalidArgument("Iceberg map writer requires an Arrow map field"); + } + const auto& map_type = assert_cast(*type); + const auto& arrow_map_type = assert_cast(*field->type()); + const auto& key_field = arrow_map_type.key_field(); + const auto& value_field = arrow_map_type.item_field(); + return write_map_column_to_target( + column, null_map, array_builder, start, end, + [&](const IColumn& nested_keys, arrow::ArrayBuilder* key_builder, int64_t nested_start, + int64_t nested_end) { + return key_serde->write_column_to_iceberg_arrow( + map_type.get_key_type(), nested_keys, nullptr, key_field, key_builder, + nested_start, nested_end, ctz); + }, + [&](const IColumn& nested_values, arrow::ArrayBuilder* value_builder, + int64_t nested_start, int64_t nested_end) { + return value_serde->write_column_to_iceberg_arrow( + map_type.get_value_type(), nested_values, nullptr, value_field, + value_builder, nested_start, nested_end, ctz); + }); +} + Status DataTypeMapSerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { diff --git a/be/src/core/data_type_serde/data_type_map_serde.h b/be/src/core/data_type_serde/data_type_map_serde.h index 478fb155afbbbf..1404963ecf3e10 100644 --- a/be/src/core/data_type_serde/data_type_map_serde.h +++ b/be/src/core/data_type_serde/data_type_map_serde.h @@ -83,6 +83,16 @@ class DataTypeMapSerDe : public DataTypeSerDe { Status write_column_to_arrow(const IColumn& column, const NullMap* null_map, arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_paimon_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_iceberg_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_orc(IColumn& column, const OrcDecodedColumnView& view) const override; diff --git a/be/src/core/data_type_serde/data_type_nullable_serde.cpp b/be/src/core/data_type_serde/data_type_nullable_serde.cpp index 115b37ea2a891c..de8c896f903848 100644 --- a/be/src/core/data_type_serde/data_type_nullable_serde.cpp +++ b/be/src/core/data_type_serde/data_type_nullable_serde.cpp @@ -31,6 +31,7 @@ #include "core/column/column_const.h" #include "core/column/column_nullable.h" #include "core/column/column_vector.h" +#include "core/data_type/data_type_nullable.h" #include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/data_type_serde.h" #include "core/data_type_serde/data_type_string_serde.h" @@ -385,6 +386,28 @@ Status DataTypeNullableSerDe::write_column_to_arrow(const IColumn& column, const start, end, ctz); } +Status DataTypeNullableSerDe::write_column_to_paimon_arrow( + const std::shared_ptr& type, const IColumn& column, const NullMap*, + const std::shared_ptr& field, arrow::ArrayBuilder* array_builder, + int64_t start, int64_t end, const cctz::time_zone& ctz) const { + const auto& nullable_type = assert_cast(*type); + const auto& column_nullable = assert_cast(column); + return nested_serde->write_column_to_paimon_arrow( + nullable_type.get_nested_type(), column_nullable.get_nested_column(), + &column_nullable.get_null_map_data(), field, array_builder, start, end, ctz); +} + +Status DataTypeNullableSerDe::write_column_to_iceberg_arrow( + const std::shared_ptr& type, const IColumn& column, const NullMap*, + const std::shared_ptr& field, arrow::ArrayBuilder* array_builder, + int64_t start, int64_t end, const cctz::time_zone& ctz) const { + const auto& nullable_type = assert_cast(*type); + const auto& column_nullable = assert_cast(column); + return nested_serde->write_column_to_iceberg_arrow( + nullable_type.get_nested_type(), column_nullable.get_nested_column(), + &column_nullable.get_null_map_data(), field, array_builder, start, end, ctz); +} + Status DataTypeNullableSerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, diff --git a/be/src/core/data_type_serde/data_type_nullable_serde.h b/be/src/core/data_type_serde/data_type_nullable_serde.h index d8f4cae2e749eb..8be1ece1e61a5f 100644 --- a/be/src/core/data_type_serde/data_type_nullable_serde.h +++ b/be/src/core/data_type_serde/data_type_nullable_serde.h @@ -87,6 +87,16 @@ class DataTypeNullableSerDe : public DataTypeSerDe { Status write_column_to_arrow(const IColumn& column, const NullMap* null_map, arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_paimon_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_iceberg_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_decoded_values(IColumn& column, diff --git a/be/src/core/data_type_serde/data_type_serde.h b/be/src/core/data_type_serde/data_type_serde.h index 67a1c38af958be..b4cb8936aa369e 100644 --- a/be/src/core/data_type_serde/data_type_serde.h +++ b/be/src/core/data_type_serde/data_type_serde.h @@ -37,6 +37,7 @@ namespace arrow { class ArrayBuilder; class Array; +class Field; } // namespace arrow namespace cctz { class time_zone; @@ -498,6 +499,23 @@ class DataTypeSerDe { virtual Status write_column_to_arrow(const IColumn& column, const NullMap* null_map, arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, const cctz::time_zone& ctz) const = 0; + // Most scalar types deliberately share their physical Arrow encoding across these protocols. + // Target-specific SerDes override the corresponding method; callers never retry another + // protocol method after an error. + virtual Status write_column_to_paimon_arrow(const std::shared_ptr&, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr&, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const { + return write_column_to_arrow(column, null_map, array_builder, start, end, ctz); + } + virtual Status write_column_to_iceberg_arrow(const std::shared_ptr&, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr&, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const { + return write_column_to_arrow(column, null_map, array_builder, start, end, ctz); + } virtual Status read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const = 0; diff --git a/be/src/core/data_type_serde/data_type_string_serde.cpp b/be/src/core/data_type_serde/data_type_string_serde.cpp index bccc79e75cd917..22987bfd720169 100644 --- a/be/src/core/data_type_serde/data_type_string_serde.cpp +++ b/be/src/core/data_type_serde/data_type_string_serde.cpp @@ -17,6 +17,9 @@ #include "core/data_type_serde/data_type_string_serde.h" +#include +#include + #include #include #include @@ -30,6 +33,7 @@ #include "core/data_type_serde/decoded_column_view.h" #include "core/data_type_serde/orc_serde_utils.h" #include "core/data_type_serde/parquet_decode_source.h" +#include "format/arrow/arrow_block_convertor.h" #include "util/jsonb_document_cast.h" #include "util/jsonb_utils.h" #include "util/jsonb_writer.h" @@ -209,6 +213,14 @@ class StringParquetConsumer final : public ParquetFixedValueConsumer, namespace { +bool is_iceberg_uuid_field(const std::shared_ptr& field) { + if (!field->HasMetadata()) { + return false; + } + const auto value = field->metadata()->Get("originalType"); + return value.ok() && value.ValueUnsafe() == "uuid"; +} + int hex_value(char c) { if (c >= '0' && c <= '9') { return c - '0'; @@ -549,6 +561,41 @@ Status DataTypeStringSerDeBase::write_column_to_arrow( } } +template +Status DataTypeStringSerDeBase::write_column_to_iceberg_arrow( + const std::shared_ptr& type, const IColumn& column, + const NullMap* null_map, const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const { + if (!is_iceberg_uuid_field(field)) { + // Keep the existing CHAR/STRING fixed-binary binding until external type mappings change. + return write_column_to_arrow(column, null_map, array_builder, start, end, ctz); + } + if (!is_string_type(type->get_primitive_type()) || + array_builder->type()->id() != arrow::Type::FIXED_SIZE_BINARY) { + return Status::InvalidArgument( + "Iceberg UUID writer is not bound for Doris type {} and Arrow field {}", + type->get_name(), field->ToString()); + } + auto& builder = assert_cast(*array_builder); + const int byte_width = + assert_cast(*builder.type()).byte_width(); + if (byte_width != 16) { + return Status::InvalidArgument("Iceberg UUID expects 16 bytes, got {}", byte_width); + } + const auto& strings = assert_cast(column); + for (int64_t row = start; row < end; ++row) { + if (null_map != nullptr && (*null_map)[row]) { + RETURN_IF_ERROR(checkArrowStatus(builder.AppendNull(), column, builder)); + continue; + } + std::array bytes; + RETURN_IF_ERROR(parse_iceberg_uuid_to_bytes(strings.get_data_at(row), &bytes)); + RETURN_IF_ERROR(checkArrowStatus(builder.Append(bytes.data()), column, builder)); + } + return Status::OK(); +} + template Status DataTypeStringSerDeBase::read_column_from_arrow( IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, diff --git a/be/src/core/data_type_serde/data_type_string_serde.h b/be/src/core/data_type_serde/data_type_string_serde.h index 6389b6be92546e..f90b27b1555452 100644 --- a/be/src/core/data_type_serde/data_type_string_serde.h +++ b/be/src/core/data_type_serde/data_type_string_serde.h @@ -200,6 +200,11 @@ class DataTypeStringSerDeBase : public DataTypeSerDe { Status write_column_to_arrow(const IColumn& column, const NullMap* null_map, arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_iceberg_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; diff --git a/be/src/core/data_type_serde/data_type_struct_serde.cpp b/be/src/core/data_type_serde/data_type_struct_serde.cpp index e8e034edd1ca89..04f378104f5d9c 100644 --- a/be/src/core/data_type_serde/data_type_struct_serde.cpp +++ b/be/src/core/data_type_serde/data_type_struct_serde.cpp @@ -20,11 +20,13 @@ #include #include "arrow/array/builder_nested.h" +#include "common/cast_set.h" #include "common/config.h" #include "common/status.h" #include "core/column/column.h" #include "core/column/column_const.h" #include "core/column/column_struct.h" +#include "core/data_type/data_type_struct.h" #include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/complex_type_deserialize_util.h" #include "core/data_type_serde/data_type_serde.h" @@ -470,6 +472,75 @@ Status DataTypeStructSerDe::write_column_to_arrow(const IColumn& column, const N return Status::OK(); } +namespace { + +template +Status write_struct_column_to_target(const IColumn& column, const NullMap* null_map, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + WriteElement&& write_element) { + auto& builder = assert_cast(*array_builder); + const auto& struct_column = assert_cast(column); + for (int64_t row = start; row < end; ++row) { + if (null_map != nullptr && (*null_map)[row]) { + RETURN_IF_ERROR(checkArrowStatus(builder.AppendNull(), struct_column, builder)); + continue; + } + RETURN_IF_ERROR(checkArrowStatus(builder.Append(), struct_column, builder)); + for (size_t element = 0; element < struct_column.tuple_size(); ++element) { + RETURN_IF_ERROR(write_element(struct_column, builder, element, row)); + } + } + return Status::OK(); +} + +} // namespace + +Status DataTypeStructSerDe::write_column_to_paimon_arrow( + const std::shared_ptr& type, const IColumn& column, + const NullMap* null_map, const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const { + const auto& struct_type = assert_cast(*type); + // Child indices are meaningful only when the target has the same structural arity. + if (field->type()->id() != arrow::Type::STRUCT || + array_builder->type()->id() != arrow::Type::STRUCT || + field->type()->num_fields() != struct_type.get_elements().size()) { + return Status::InvalidArgument("Paimon struct writer requires matching Arrow fields"); + } + return write_struct_column_to_target( + column, null_map, array_builder, start, end, + [&](const ColumnStruct& struct_column, arrow::StructBuilder& builder, size_t element, + int64_t row) { + return elem_serdes_ptrs[element]->write_column_to_paimon_arrow( + struct_type.get_element(element), struct_column.get_column(element), + nullptr, field->type()->field(cast_set(element)), + builder.field_builder(cast_set(element)), row, row + 1, ctz); + }); +} + +Status DataTypeStructSerDe::write_column_to_iceberg_arrow( + const std::shared_ptr& type, const IColumn& column, + const NullMap* null_map, const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const { + const auto& struct_type = assert_cast(*type); + // Child indices are meaningful only when the target has the same structural arity. + if (field->type()->id() != arrow::Type::STRUCT || + array_builder->type()->id() != arrow::Type::STRUCT || + field->type()->num_fields() != struct_type.get_elements().size()) { + return Status::InvalidArgument("Iceberg struct writer requires matching Arrow fields"); + } + return write_struct_column_to_target( + column, null_map, array_builder, start, end, + [&](const ColumnStruct& struct_column, arrow::StructBuilder& builder, size_t element, + int64_t row) { + return elem_serdes_ptrs[element]->write_column_to_iceberg_arrow( + struct_type.get_element(element), struct_column.get_column(element), + nullptr, field->type()->field(cast_set(element)), + builder.field_builder(cast_set(element)), row, row + 1, ctz); + }); +} + Status DataTypeStructSerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { diff --git a/be/src/core/data_type_serde/data_type_struct_serde.h b/be/src/core/data_type_serde/data_type_struct_serde.h index dd3ff38f1609d6..a6a810940bc756 100644 --- a/be/src/core/data_type_serde/data_type_struct_serde.h +++ b/be/src/core/data_type_serde/data_type_struct_serde.h @@ -84,6 +84,16 @@ class DataTypeStructSerDe : public DataTypeSerDe { Status write_column_to_arrow(const IColumn& column, const NullMap* null_map, arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_paimon_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_iceberg_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_orc(IColumn& column, const OrcDecodedColumnView& view) const override; diff --git a/be/src/core/data_type_serde/data_type_variant_v2_serde.cpp b/be/src/core/data_type_serde/data_type_variant_v2_serde.cpp index b970c49ac31b6a..0e0d7a58d2f091 100644 --- a/be/src/core/data_type_serde/data_type_variant_v2_serde.cpp +++ b/be/src/core/data_type_serde/data_type_variant_v2_serde.cpp @@ -19,6 +19,7 @@ #include #include +#include #include #include @@ -183,187 +184,6 @@ void preflight_json(const IColumn& column, size_t start, size_t end, }); } -void validate_paimon_variant_primitive(VariantPrimitiveId primitive_id) { - switch (primitive_id) { - case VariantPrimitiveId::NULL_VALUE: - case VariantPrimitiveId::TRUE_VALUE: - case VariantPrimitiveId::FALSE_VALUE: - case VariantPrimitiveId::INT8: - case VariantPrimitiveId::INT16: - case VariantPrimitiveId::INT32: - case VariantPrimitiveId::INT64: - case VariantPrimitiveId::DOUBLE: - case VariantPrimitiveId::DECIMAL4: - case VariantPrimitiveId::DECIMAL8: - case VariantPrimitiveId::DECIMAL16: - case VariantPrimitiveId::DATE: - case VariantPrimitiveId::TIMESTAMP_MICROS: - case VariantPrimitiveId::TIMESTAMP_NTZ_MICROS: - case VariantPrimitiveId::FLOAT: - case VariantPrimitiveId::BINARY: - case VariantPrimitiveId::STRING: - case VariantPrimitiveId::UUID: - return; - case VariantPrimitiveId::TIME_NTZ_MICROS: - case VariantPrimitiveId::TIMESTAMP_NANOS: - case VariantPrimitiveId::TIMESTAMP_NTZ_NANOS: - throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, - "Paimon does not support Variant primitive id {}", - static_cast(primitive_id)); - } - throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, - "Paimon does not support unknown Variant primitive id {}", - static_cast(primitive_id)); -} - -void validate_paimon_variant_value(VariantRef value, uint32_t depth = 0) { - if (depth > VARIANT_MAX_NESTING_DEPTH) { - throw Exception(ErrorCode::CORRUPTION, "Variant value exceeds maximum nesting depth {}", - VARIANT_MAX_NESTING_DEPTH); - } - const size_t encoded_size = value.value_size(); - if (encoded_size != value.value.size) { - throw Exception(ErrorCode::CORRUPTION, - "Variant value has {} trailing bytes after the encoded value", - value.value.size - encoded_size); - } - - switch (value.basic_type()) { - case VariantBasicType::PRIMITIVE: - validate_paimon_variant_primitive(value.primitive_id()); - return; - case VariantBasicType::SHORT_STRING: - return; - case VariantBasicType::OBJECT: - for (uint32_t i = 0; i < value.num_elements(); ++i) { - uint32_t field_id = 0; - VariantRef child = value.object_value_at(i, &field_id); - value.metadata.key_at(field_id); - validate_paimon_variant_value(child, depth + 1); - } - return; - case VariantBasicType::ARRAY: - for (uint32_t i = 0; i < value.num_elements(); ++i) { - validate_paimon_variant_value(value.array_at(i), depth + 1); - } - return; - } -} - -void require_variant_arrow_status(const arrow::Status& status) { - if (!status.ok()) { - throw Exception(ErrorCode::INTERNAL_ERROR, "Variant V2 Arrow append failed: {}", - status.ToString()); - } -} - -Status write_binary_variant_arrow(const IColumn& column, const NullMap* null_map, - arrow::StructBuilder& builder, size_t start, size_t end) { - // StructBuilder::type() returns a shared_ptr by value. Keep that owner alive while using the - // cast reference; otherwise the reference would dangle as soon as the temporary is destroyed. - const auto builder_type = builder.type(); - const auto& struct_type = assert_cast(*builder_type); - if (struct_type.num_fields() != 2 || struct_type.field(0)->name() != "value" || - struct_type.field(1)->name() != "metadata" || - struct_type.field(0)->type()->id() != arrow::Type::BINARY || - struct_type.field(1)->type()->id() != arrow::Type::BINARY) { - return Status::InvalidArgument( - "Binary Variant V2 Arrow type must be " - "struct, got {}", - struct_type.ToString()); - } - auto* value_builder = dynamic_cast(builder.field_builder(0)); - auto* metadata_builder = dynamic_cast(builder.field_builder(1)); - if (value_builder == nullptr || metadata_builder == nullptr) { - return Status::InvalidArgument("Binary Variant V2 Arrow child builders must be binary"); - } - - // GenericVariant assumes its input is valid, and Paimon's unshredded writer copies these two - // buffers without inspecting them. Validate once at the Doris-to-Paimon boundary so a write - // cannot commit bytes which Paimon is unable to read later. - const auto outer_nulls = forced_nulls(null_map); - visit_variant_v2_values( - column, start, end, outer_nulls, - [&](size_t) { require_variant_arrow_status(builder.AppendNull()); }, - [&](size_t row, VariantRef value) { - try { - constexpr size_t PAIMON_VARIANT_SIZE_LIMIT = 128 * 1024 * 1024; - if (value.value.size > PAIMON_VARIANT_SIZE_LIMIT || - value.metadata.size > PAIMON_VARIANT_SIZE_LIMIT) { - throw Exception(ErrorCode::INVALID_ARGUMENT, - "exceeds the 128 MiB value/metadata limit"); - } - value.metadata.validate(); - validate_paimon_variant_value(value); - } catch (const Exception& e) { - throw Exception(e.code(), "Paimon Variant V2 row {} is incompatible: {}", row, - e.what()); - } - require_variant_arrow_status(builder.Append()); - require_variant_arrow_status( - value_builder->Append(reinterpret_cast(value.value.data), - cast_set(value.value.size))); - require_variant_arrow_status(metadata_builder->Append( - reinterpret_cast(value.metadata.data), - cast_set(value.metadata.size))); - }); - return Status::OK(); -} - -Status write_arrow_variant_storage(const IColumn& column, const NullMap* null_map, - arrow::StructBuilder& builder, size_t start, size_t end) { - const auto builder_type = builder.type(); - const auto& struct_type = assert_cast(*builder_type); - if (struct_type.num_fields() != 2 || struct_type.field(0)->name() != "metadata" || - struct_type.field(1)->name() != "value" || - struct_type.field(0)->type()->id() != arrow::Type::BINARY || - struct_type.field(1)->type()->id() != arrow::Type::BINARY) { - return Status::InvalidArgument( - "Iceberg Variant Arrow storage must be " - "struct, got {}", - struct_type.ToString()); - } - auto* metadata_builder = dynamic_cast(builder.field_builder(0)); - auto* value_builder = dynamic_cast(builder.field_builder(1)); - if (metadata_builder == nullptr || value_builder == nullptr) { - return Status::InvalidArgument("Iceberg Variant Arrow storage fields must both be binary"); - } - - Status status = Status::OK(); - visit_variant_v2_values( - column, start, end, forced_nulls(null_map), - [&](size_t) { - if (status.ok()) { - status = checkArrowStatus(builder.AppendNull(), column, builder); - } - }, - [&](size_t, VariantRef value) { - if (!status.ok()) { - return; - } - if (value.metadata.size > std::numeric_limits::max() || - value.value.size > std::numeric_limits::max()) { - status = Status::InvalidArgument( - "Iceberg Variant metadata/value exceeds Arrow binary size limit"); - return; - } - status = checkArrowStatus(builder.Append(), column, builder); - if (status.ok()) { - status = checkArrowStatus( - metadata_builder->Append(value.metadata.data, - static_cast(value.metadata.size)), - column, *metadata_builder); - } - if (status.ok()) { - status = checkArrowStatus( - value_builder->Append(value.value.data, - static_cast(value.value.size)), - column, *value_builder); - } - }); - return status; -} - } // namespace DataTypeVariantV2SerDe::DataTypeVariantV2SerDe(int nesting_level) : DataTypeSerDe(nesting_level) {} @@ -670,6 +490,196 @@ Status write_arrow(const IColumn& column, const NullMap* null_map, Builder& buil return status; } +constexpr size_t PAIMON_VARIANT_SIZE_LIMIT = 128 * 1024 * 1024; + +void validate_paimon_variant_primitive(VariantPrimitiveId primitive_id) { + switch (primitive_id) { + case VariantPrimitiveId::NULL_VALUE: + case VariantPrimitiveId::TRUE_VALUE: + case VariantPrimitiveId::FALSE_VALUE: + case VariantPrimitiveId::INT8: + case VariantPrimitiveId::INT16: + case VariantPrimitiveId::INT32: + case VariantPrimitiveId::INT64: + case VariantPrimitiveId::DOUBLE: + case VariantPrimitiveId::DECIMAL4: + case VariantPrimitiveId::DECIMAL8: + case VariantPrimitiveId::DECIMAL16: + case VariantPrimitiveId::DATE: + case VariantPrimitiveId::TIMESTAMP_MICROS: + case VariantPrimitiveId::TIMESTAMP_NTZ_MICROS: + case VariantPrimitiveId::FLOAT: + case VariantPrimitiveId::BINARY: + case VariantPrimitiveId::STRING: + case VariantPrimitiveId::UUID: + return; + case VariantPrimitiveId::TIME_NTZ_MICROS: + case VariantPrimitiveId::TIMESTAMP_NANOS: + case VariantPrimitiveId::TIMESTAMP_NTZ_NANOS: + throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, + "Paimon does not support Variant primitive id {}", + static_cast(primitive_id)); + } + throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, + "Paimon does not support unknown Variant primitive id {}", + static_cast(primitive_id)); +} + +void validate_paimon_variant_value(VariantRef value, uint32_t depth = 0) { + if (depth > VARIANT_MAX_NESTING_DEPTH) { + throw Exception(ErrorCode::CORRUPTION, "Variant value exceeds maximum nesting depth {}", + VARIANT_MAX_NESTING_DEPTH); + } + const size_t encoded_size = value.value_size(); + if (encoded_size != value.value.size) { + throw Exception(ErrorCode::CORRUPTION, + "Variant value has {} trailing bytes after the encoded value", + value.value.size - encoded_size); + } + + switch (value.basic_type()) { + case VariantBasicType::PRIMITIVE: + validate_paimon_variant_primitive(value.primitive_id()); + return; + case VariantBasicType::SHORT_STRING: + return; + case VariantBasicType::OBJECT: + for (uint32_t index = 0; index < value.num_elements(); ++index) { + uint32_t field_id = 0; + const VariantRef child = value.object_value_at(index, &field_id); + value.metadata.key_at(field_id); + validate_paimon_variant_value(child, depth + 1); + } + return; + case VariantBasicType::ARRAY: + for (uint32_t index = 0; index < value.num_elements(); ++index) { + validate_paimon_variant_value(value.array_at(index), depth + 1); + } + return; + } +} + +Status write_paimon_variant(const IColumn& column, const NullMap* null_map, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end) { + if (start < 0 || end < start) { + return Status::InvalidArgument("Invalid Paimon Variant row range [{}, {})", start, end); + } + if (array_builder->type()->id() != arrow::Type::STRUCT) { + return Status::InvalidArgument("Paimon Variant writer requires a struct builder, got {}", + array_builder->type()->ToString()); + } + auto& builder = assert_cast(*array_builder); + const auto type = std::static_pointer_cast(builder.type()); + if (type->num_fields() != 2 || type->field(0)->name() != "value" || + type->field(1)->name() != "metadata" || + type->field(0)->type()->id() != arrow::Type::BINARY || + type->field(1)->type()->id() != arrow::Type::BINARY) { + return Status::InvalidArgument( + "Paimon Variant writer requires struct, got {}", + type->ToString()); + } + auto& value_builder = assert_cast(*builder.field_builder(0)); + auto& metadata_builder = assert_cast(*builder.field_builder(1)); + Status status = Status::OK(); + visit_variant_v2_values( + column, start, end, forced_nulls(null_map), + [&](size_t) { + if (status.ok()) { + status = checkArrowStatus(builder.AppendNull(), column, builder); + } + }, + [&](size_t row, VariantRef value) { + if (!status.ok()) { + return; + } + try { + if (value.value.size > PAIMON_VARIANT_SIZE_LIMIT || + value.metadata.size > PAIMON_VARIANT_SIZE_LIMIT) { + throw Exception(ErrorCode::INVALID_ARGUMENT, + "exceeds the 128 MiB value/metadata limit"); + } + value.metadata.validate(); + validate_paimon_variant_value(value); + } catch (const Exception& e) { + status = Status::Error(e.code(), + "Paimon Variant V2 row {} is incompatible: {}", + row, e.what()); + return; + } + status = checkArrowStatus(builder.Append(), column, builder); + if (status.ok()) { + status = checkArrowStatus( + value_builder.Append( + reinterpret_cast(value.value.data), + cast_set(value.value.size)), + column, value_builder); + } + if (status.ok()) { + status = checkArrowStatus( + metadata_builder.Append( + reinterpret_cast(value.metadata.data), + cast_set(value.metadata.size)), + column, metadata_builder); + } + }); + return status; +} + +Status write_iceberg_variant(const IColumn& column, const NullMap* null_map, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end) { + if (start < 0 || end < start) { + return Status::InvalidArgument("Invalid Iceberg Variant row range [{}, {})", start, end); + } + if (array_builder->type()->id() != arrow::Type::STRUCT) { + return Status::InvalidArgument("Iceberg Variant writer requires a struct builder, got {}", + array_builder->type()->ToString()); + } + auto& builder = assert_cast(*array_builder); + const auto type = std::static_pointer_cast(builder.type()); + if (type->num_fields() != 2 || type->field(0)->name() != "metadata" || + type->field(1)->name() != "value" || type->field(0)->type()->id() != arrow::Type::BINARY || + type->field(1)->type()->id() != arrow::Type::BINARY) { + return Status::InvalidArgument( + "Iceberg Variant writer requires struct, got {}", + type->ToString()); + } + auto& metadata_builder = assert_cast(*builder.field_builder(0)); + auto& value_builder = assert_cast(*builder.field_builder(1)); + Status status = Status::OK(); + visit_variant_v2_values( + column, start, end, forced_nulls(null_map), + [&](size_t) { + if (status.ok()) { + status = checkArrowStatus(builder.AppendNull(), column, builder); + } + }, + [&](size_t, VariantRef value) { + if (!status.ok()) { + return; + } + if (value.metadata.size > std::numeric_limits::max() || + value.value.size > std::numeric_limits::max()) { + status = Status::InvalidArgument( + "Iceberg Variant metadata/value exceeds Arrow binary size limit"); + return; + } + status = checkArrowStatus(builder.Append(), column, builder); + if (status.ok()) { + status = checkArrowStatus( + metadata_builder.Append(value.metadata.data, + cast_set(value.metadata.size)), + column, metadata_builder); + } + if (status.ok()) { + status = checkArrowStatus( + value_builder.Append(value.value.data, + cast_set(value.value.size)), + column, value_builder); + } + }); + return status; +} + } // namespace void DataTypeVariantV2SerDe::to_string(const IColumn& column, size_t row_num, BufferWritable& bw, @@ -732,22 +742,26 @@ Status DataTypeVariantV2SerDe::write_column_to_arrow(const IColumn& column, cons assert_cast(*array_builder), first, last, options); } - if (array_builder->type()->id() == arrow::Type::STRUCT) { - auto& struct_builder = assert_cast(*array_builder); - const auto builder_type = struct_builder.type(); - const auto& struct_type = assert_cast(*builder_type); - if (struct_type.num_fields() == 2 && struct_type.field(0)->name() == "metadata" && - struct_type.field(1)->name() == "value") { - return write_arrow_variant_storage(column, null_map, struct_builder, first, last); - } - return write_binary_variant_arrow(column, null_map, struct_builder, first, last); - } return Status::InvalidArgument("Unsupported arrow type for variant column: {}", array_builder->type()->name()); }); return Status::OK(); } +Status DataTypeVariantV2SerDe::write_column_to_paimon_arrow( + const std::shared_ptr&, const IColumn& column, const NullMap* null_map, + const std::shared_ptr&, arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone&) const { + return write_paimon_variant(column, null_map, array_builder, start, end); +} + +Status DataTypeVariantV2SerDe::write_column_to_iceberg_arrow( + const std::shared_ptr&, const IColumn& column, const NullMap* null_map, + const std::shared_ptr&, arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone&) const { + return write_iceberg_variant(column, null_map, array_builder, start, end); +} + Status DataTypeVariantV2SerDe::write_column_to_orc(const std::string&, const IColumn& column, const NullMap* null_map, orc::ColumnVectorBatch* orc_col_batch, diff --git a/be/src/core/data_type_serde/data_type_variant_v2_serde.h b/be/src/core/data_type_serde/data_type_variant_v2_serde.h index 23a3dcaa941406..ff2565784b6e6e 100644 --- a/be/src/core/data_type_serde/data_type_variant_v2_serde.h +++ b/be/src/core/data_type_serde/data_type_variant_v2_serde.h @@ -56,6 +56,16 @@ class DataTypeVariantV2SerDe final : public DataTypeSerDe { Status write_column_to_arrow(const IColumn& column, const NullMap* null_map, arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_paimon_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; + Status write_column_to_iceberg_arrow(const std::shared_ptr& type, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const override; Status read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; Status write_column_to_mysql_binary(const IColumn& column, MysqlRowBinaryBuffer& row_buffer, diff --git a/be/src/exec/operator/memory_scratch_sink_operator.cpp b/be/src/exec/operator/memory_scratch_sink_operator.cpp index 916abc461c7a7f..7b3bf3eed0b8b7 100644 --- a/be/src/exec/operator/memory_scratch_sink_operator.cpp +++ b/be/src/exec/operator/memory_scratch_sink_operator.cpp @@ -103,17 +103,16 @@ Status MemoryScratchSinkOperatorX::sink_impl(RuntimeState* state, Block* input_b RETURN_IF_ERROR(VExprContext::get_output_block_after_execute_exprs( local_state._output_vexpr_ctxs, *input_block, &block)); } - std::shared_ptr block_arrow_schema; + DorisArrowBlockConvertor converter(block, state->timezone(), _timezone_obj, + /*datetime_naive=*/true); { SCOPED_TIMER(local_state._get_arrow_schema_timer); // After expr executed, use recaculated schema as final schema - RETURN_IF_ERROR(get_arrow_schema_from_block(block, &block_arrow_schema, state->timezone(), - /*datetime_naive=*/true)); + RETURN_IF_ERROR(converter.init()); } { SCOPED_TIMER(local_state._convert_block_to_arrow_batch_timer); - RETURN_IF_ERROR(convert_to_arrow_batch( - block, block_arrow_schema, arrow::default_memory_pool(), &result, _timezone_obj)); + RETURN_IF_ERROR(converter.convert_to_arrow(block, arrow::default_memory_pool(), &result)); } local_state._queue->blocking_put(result); if (local_state._queue->size() > config::max_memory_sink_batch_count) { diff --git a/be/src/exec/sink/writer/iceberg/viceberg_delete_file_writer.cpp b/be/src/exec/sink/writer/iceberg/viceberg_delete_file_writer.cpp index 0225d3a8457856..8ad7751e1ed19d 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_delete_file_writer.cpp +++ b/be/src/exec/sink/writer/iceberg/viceberg_delete_file_writer.cpp @@ -21,8 +21,8 @@ #include "format/table/iceberg/schema.h" #include "format/table/iceberg/types.h" +#include "format/transformer/viceberg_parquet_writer.h" #include "format/transformer/vorc_transformer.h" -#include "format/transformer/vparquet_transformer.h" #include "io/file_factory.h" #include "runtime/runtime_state.h" @@ -104,9 +104,9 @@ Status VIcebergDeleteFileWriter::open(RuntimeState* state, RuntimeProfile* profi ParquetFileOptions parquet_options = {parquet_compression_type, TParquetVersion::PARQUET_1_0, false, false}; - _file_format_transformer.reset(new VParquetTransformer( + _file_format_transformer.reset(new VIcebergParquetWriter( state, _file_writer.get(), output_exprs, column_names, false, parquet_options, - nullptr, _position_delete_schema.get())); + nullptr, *_position_delete_schema)); return _file_format_transformer->open(); } case TFileFormatType::FORMAT_ORC: { diff --git a/be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp b/be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp index ba7644daec751f..e1d2de9c439652 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp +++ b/be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp @@ -23,8 +23,8 @@ #include "core/block/materialize_block.h" #include "core/column/column_map.h" #include "format/table/iceberg/schema.h" +#include "format/transformer/viceberg_parquet_writer.h" #include "format/transformer/vorc_transformer.h" -#include "format/transformer/vparquet_transformer.h" #include "io/file_factory.h" #include "runtime/runtime_state.h" @@ -100,9 +100,9 @@ Status VIcebergPartitionWriter::open(RuntimeState* state, RuntimeProfile* profil .parquet_version = TParquetVersion::PARQUET_1_0, .parquet_disable_dictionary = false, .enable_int96_timestamps = false}; - _file_format_transformer = std::make_unique( + _file_format_transformer = std::make_unique( state, _file_writer.get(), _write_output_expr_ctxs, _write_column_names, false, - parquet_options, _iceberg_schema_json, &_schema); + parquet_options, _iceberg_schema_json, _schema); return _file_format_transformer->open(); } case TFileFormatType::FORMAT_ORC: { @@ -178,7 +178,7 @@ Status VIcebergPartitionWriter::_build_iceberg_commit_data(TIcebergCommitData* c } if (_file_format_type == TFileFormatType::FORMAT_PARQUET) { TIcebergColumnStats column_stats; - RETURN_IF_ERROR(static_cast(_file_format_transformer.get()) + RETURN_IF_ERROR(static_cast(_file_format_transformer.get()) ->collect_file_statistics_after_close(&column_stats)); commit_data->__set_column_stats(column_stats); } else if (_file_format_type == TFileFormatType::FORMAT_ORC) { diff --git a/be/src/exec/sink/writer/paimon/jni_paimon_write_backend.cpp b/be/src/exec/sink/writer/paimon/jni_paimon_write_backend.cpp index d46fb1d99f691f..923143051f4329 100644 --- a/be/src/exec/sink/writer/paimon/jni_paimon_write_backend.cpp +++ b/be/src/exec/sink/writer/paimon/jni_paimon_write_backend.cpp @@ -17,10 +17,7 @@ #include "exec/sink/writer/paimon/jni_paimon_write_backend.h" -#include #include -#include -#include #include #include @@ -36,6 +33,7 @@ #include "exec/sink/writer/paimon/paimon_jni_memory_manager.h" #include "exec/spill/spill_file_manager.h" #include "format/arrow/arrow_block_convertor.h" +#include "format/table/paimon/paimon_arrow_block_convertor.h" #include "runtime/exec_env.h" #include "runtime/query_context.h" #include "runtime/runtime_state.h" @@ -215,7 +213,7 @@ JniPaimonWriteBackend::~JniPaimonWriteBackend() { Status JniPaimonWriteBackend::close() { if (_jni_writer_obj == nullptr && _jni_writer_cls == nullptr) { _memory_manager.reset(); - _arrow_schema.reset(); + _serialized_arrow_schema.clear(); _spill_session.reset(); _opened = false; return Status::OK(); @@ -236,7 +234,7 @@ Status JniPaimonWriteBackend::close() { _memory_manager.reset(); _spill_session.reset(); } - _arrow_schema.reset(); + _serialized_arrow_schema.clear(); _opened = false; return env_status; } @@ -272,7 +270,7 @@ Status JniPaimonWriteBackend::close() { // Retain ownership until process exit and fence subsequent writer admission. retain_resources_after_failed_close(std::move(_memory_manager), std::move(_spill_session)); } - _arrow_schema.reset(); + _serialized_arrow_schema.clear(); _opened = false; return close_status; } @@ -288,7 +286,7 @@ Status JniPaimonWriteBackend::_check_jni_exception(JNIEnv* env, const std::strin } static Status _get_paimon_arrow_schema(JNIEnv* env, jobject writer, jmethodID get_schema_id, - std::shared_ptr* schema) { + std::string* serialized_schema) { auto schema_bytes = static_cast(env->CallObjectMethod(writer, get_schema_id)); RETURN_IF_ERROR(Jni::Env::GetJniExceptionMsg( env, false, "JNI exception in PaimonJniWriter.getArrowSchema: ")); @@ -301,21 +299,13 @@ static Status _get_paimon_arrow_schema(JNIEnv* env, jobject writer, jmethodID ge env->DeleteLocalRef(schema_bytes); return Status::InternalError("PaimonJniWriter.getArrowSchema returned empty data"); } - std::string serialized_schema(static_cast(size), '\0'); + serialized_schema->assign(static_cast(size), '\0'); env->GetByteArrayRegion(schema_bytes, 0, size, - reinterpret_cast(serialized_schema.data())); + reinterpret_cast(serialized_schema->data())); env->DeleteLocalRef(schema_bytes); RETURN_IF_ERROR(Jni::Env::GetJniExceptionMsg( env, false, "JNI exception while reading Paimon Arrow schema: ")); - auto input = std::make_shared( - arrow::Buffer::FromString(std::move(serialized_schema))); - auto reader_result = arrow::ipc::RecordBatchStreamReader::Open(input); - if (!reader_result.ok()) { - return Status::InternalError("Failed to deserialize Paimon Arrow schema: {}", - reader_result.status().ToString()); - } - *schema = reader_result.ValueOrDie()->schema(); return Status::OK(); } Status JniPaimonWriteBackend::open(const TPaimonTableSink& sink, RuntimeState* state, @@ -325,7 +315,8 @@ Status JniPaimonWriteBackend::open(const TPaimonTableSink& sink, RuntimeState* s "Paimon JNI writes are disabled on this BE because a previous Java writer close " "could not be confirmed; restart the BE to reclaim retained native memory safely"); } - _arrow_schema.reset(); + _serialized_arrow_schema.clear(); + _timezone = state->timezone_obj(); DORIS_CHECK(sink.__isset.column_names); DORIS_CHECK(sink.__isset.write_mode); DORIS_CHECK(sink.__isset.serialized_table); @@ -421,7 +412,8 @@ Status JniPaimonWriteBackend::open(const TPaimonTableSink& sink, RuntimeState* s Status st = _check_jni_exception(env, "open PaimonJniWriter"); if (st.ok()) { - st = _get_paimon_arrow_schema(env, _jni_writer_obj, get_arrow_schema_id, &_arrow_schema); + st = _get_paimon_arrow_schema(env, _jni_writer_obj, get_arrow_schema_id, + &_serialized_arrow_schema); } if (st.ok()) { _opened = true; @@ -438,38 +430,33 @@ Status JniPaimonWriteBackend::open(const TPaimonTableSink& sink, RuntimeState* s Status JniPaimonWriteBackend::create_writer( // NOLINT(readability-make-member-function-const) std::unique_ptr* writer) { DORIS_CHECK(_opened); - DORIS_CHECK(_arrow_schema != nullptr); - *writer = std::make_unique(_jni_writer_obj, _write_id, _prepare_commit_id, - _abort_id, _arrow_schema); + DORIS_CHECK(!_serialized_arrow_schema.empty()); + auto paimon_writer = + std::make_unique(_jni_writer_obj, _write_id, _prepare_commit_id, + _abort_id, _serialized_arrow_schema, _timezone); + RETURN_IF_ERROR(paimon_writer->init()); + *writer = std::move(paimon_writer); return Status::OK(); } JniPaimonWriter::JniPaimonWriter(jobject jni_writer_obj, jmethodID write_id, jmethodID prepare_commit_id, jmethodID abort_id, - std::shared_ptr arrow_schema) + std::string serialized_arrow_schema, + const cctz::time_zone& timezone) : _jni_writer_obj(jni_writer_obj), _write_id(write_id), _prepare_commit_id(prepare_commit_id), _abort_id(abort_id), - _arrow_schema(std::move(arrow_schema)) {} + _arrow_block_convertor(std::move(serialized_arrow_schema), timezone) {} Status JniPaimonWriter::write(RuntimeState* state, Block& block) { if (block.rows() == 0) { return Status::OK(); } - if (_arrow_schema == nullptr || _arrow_schema->num_fields() != block.columns()) { - return Status::InvalidArgument( - "Paimon Arrow schema column count does not match Doris Block: schema={}, block={}", - _arrow_schema == nullptr ? 0 : _arrow_schema->num_fields(), block.columns()); - } - - // The schema comes from the pinned Paimon table, so timestamp timezone, nested nullability and - // Variant layout are fixed before the first write. Arrow builders remain on the Doris side and - // are charged to the current query's MemTracker through ArrowMemoryPool. + // The converter owns the pinned table schema; builders are charged to this writer's pool. std::shared_ptr record_batch; - RETURN_IF_ERROR(convert_to_arrow_batch(block, _arrow_schema, &_arrow_pool, &record_batch, - state->timezone_obj())); + RETURN_IF_ERROR(_arrow_block_convertor.convert_to_arrow(block, &_arrow_pool, &record_batch)); ArrowArray c_array {}; ArrowSchema c_schema {}; diff --git a/be/src/exec/sink/writer/paimon/jni_paimon_write_backend.h b/be/src/exec/sink/writer/paimon/jni_paimon_write_backend.h index d5a5de3b9ea1bb..ebf7ae812675a2 100644 --- a/be/src/exec/sink/writer/paimon/jni_paimon_write_backend.h +++ b/be/src/exec/sink/writer/paimon/jni_paimon_write_backend.h @@ -27,6 +27,7 @@ #include "exec/sink/writer/paimon/paimon_jni_memory_manager.h" #include "exec/sink/writer/paimon/paimon_write_backend.h" #include "format/parquet/arrow_memory_pool.h" +#include "format/table/paimon/paimon_arrow_block_convertor.h" #include "runtime/runtime_profile.h" namespace arrow { @@ -82,7 +83,8 @@ class JniPaimonWriteBackend final : public IPaimonWriteBackend { jmethodID _close_id = nullptr; std::unique_ptr _memory_manager; - std::shared_ptr _arrow_schema; + std::string _serialized_arrow_schema; + cctz::time_zone _timezone; std::unique_ptr _spill_session; RuntimeProfile::Counter* _native_page_memory_limit = nullptr; RuntimeProfile::Counter* _native_page_memory_peak = nullptr; @@ -97,7 +99,10 @@ class JniPaimonWriteBackend final : public IPaimonWriteBackend { class JniPaimonWriter final : public IPaimonWriter { public: JniPaimonWriter(jobject jni_writer_obj, jmethodID write_id, jmethodID prepare_commit_id, - jmethodID abort_id, std::shared_ptr arrow_schema); + jmethodID abort_id, std::string serialized_arrow_schema, + const cctz::time_zone& timezone); + + Status init() { return _arrow_block_convertor.init(); } Status write(RuntimeState* state, Block& block) override; Status prepare_commit(std::vector& messages) override; @@ -112,7 +117,7 @@ class JniPaimonWriter final : public IPaimonWriter { // Arrow resources owned by this writer adapter. ArrowMemoryPool<> _arrow_pool; - std::shared_ptr _arrow_schema; + paimon::PaimonArrowBlockConvertor _arrow_block_convertor; }; } // namespace doris diff --git a/be/src/exec/sink/writer/vfile_result_writer.cpp b/be/src/exec/sink/writer/vfile_result_writer.cpp index 0eff1fd42f830c..ce9e73d2eba5d5 100644 --- a/be/src/exec/sink/writer/vfile_result_writer.cpp +++ b/be/src/exec/sink/writer/vfile_result_writer.cpp @@ -44,7 +44,7 @@ #include "exprs/vexpr_context.h" #include "format/transformer/vcsv_transformer.h" #include "format/transformer/vorc_transformer.h" -#include "format/transformer/vparquet_transformer.h" +#include "format/transformer/vparquet_writer.h" #include "io/file_factory.h" #include "io/fs/broker_file_system.h" #include "io/fs/file_system.h" @@ -137,7 +137,7 @@ Status VFileResultWriter::_create_file_writer(const std::string& file_name) { _file_opts->with_bom, _file_opts->compression_type)); break; case TFileFormatType::FORMAT_PARQUET: - _vfile_writer.reset(new VParquetTransformer( + _vfile_writer.reset(new VParquetWriter( _state, _file_writer_impl.get(), _vec_output_expr_ctxs, _file_opts->parquet_schemas, _output_object_data, {_file_opts->parquet_commpression_type, _file_opts->parquet_version, diff --git a/be/src/exec/sink/writer/vhive_partition_writer.cpp b/be/src/exec/sink/writer/vhive_partition_writer.cpp index 8331efac54bd47..4763f5f6be9aae 100644 --- a/be/src/exec/sink/writer/vhive_partition_writer.cpp +++ b/be/src/exec/sink/writer/vhive_partition_writer.cpp @@ -22,8 +22,8 @@ #include "core/block/materialize_block.h" #include "core/column/column_map.h" #include "format/transformer/vcsv_transformer.h" +#include "format/transformer/vhive_parquet_writer.h" #include "format/transformer/vorc_transformer.h" -#include "format/transformer/vparquet_transformer.h" #include "io/file_factory.h" #include "io/fs/s3_file_writer.h" #include "runtime/runtime_state.h" @@ -97,7 +97,7 @@ Status VHivePartitionWriter::open(RuntimeState* state, RuntimeProfile* operator_ // changing the default Hive parquet timestamp encoding to standard logical types. ParquetFileOptions parquet_options = {parquet_compression_type, TParquetVersion::PARQUET_1_0, false, true}; - _file_format_transformer = std::make_unique( + _file_format_transformer = std::make_unique( state, _file_writer.get(), _write_output_expr_ctxs, _write_column_names, false, parquet_options); return _file_format_transformer->open(); diff --git a/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp b/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp index ac37602f638393..495daee2018010 100644 --- a/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp +++ b/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp @@ -59,17 +59,17 @@ Status AggregatePythonUDAFData::add(int64_t place_id, const IColumn** columns, ColumnWithTypeAndName(columns[i]->get_ptr(), argument_types[i], std::to_string(i))); } - std::shared_ptr schema; - RETURN_IF_ERROR( - get_arrow_schema_from_block(input_block, &schema, TimezoneUtils::default_time_zone)); cctz::time_zone timezone_obj; TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, timezone_obj); + PythonArrowBlockConvertor converter(input_block, TimezoneUtils::default_time_zone, + timezone_obj); + RETURN_IF_ERROR(converter.init()); std::shared_ptr batch; // Zero-copy: convert only the specified range - RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema, arrow::default_memory_pool(), - &batch, timezone_obj, row_num_start, row_num_end)); - // Send the batch (already sliced in convert_to_arrow_batch) + RETURN_IF_ERROR(converter.convert_to_arrow(input_block, arrow::default_memory_pool(), &batch, + row_num_start, row_num_end)); + // Send the batch (already sliced by the converter) // Single place mode: no places column needed RETURN_IF_ERROR(client->accumulate(place_id, true, *batch, 0, batch->num_rows())); return Status::OK(); @@ -103,17 +103,17 @@ Status AggregatePythonUDAFData::add_batch(AggregateDataPtr* places, size_t place DataTypeFactory::instance().create_data_type(PrimitiveType::TYPE_BIGINT, false); input_block.insert(ColumnWithTypeAndName(std::move(places_col), places_type, "places")); - std::shared_ptr schema; - RETURN_IF_ERROR( - get_arrow_schema_from_block(input_block, &schema, TimezoneUtils::default_time_zone)); cctz::time_zone timezone_obj; TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, timezone_obj); + PythonArrowBlockConvertor converter(input_block, TimezoneUtils::default_time_zone, + timezone_obj); + RETURN_IF_ERROR(converter.init()); std::shared_ptr batch; // Zero-copy: convert only the [start, end) range // This slice includes the places column automatically - RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema, arrow::default_memory_pool(), - &batch, timezone_obj, start, end)); + RETURN_IF_ERROR(converter.convert_to_arrow(input_block, arrow::default_memory_pool(), &batch, + start, end)); // Send entire batch (already contains places column) to Python // place_id=0 is ignored when is_single_place=false RETURN_IF_ERROR(client->accumulate(0, false, *batch, 0, slice_rows)); @@ -174,7 +174,8 @@ Status AggregatePythonUDAFData::get(IColumn& to, const DataTypePtr& result_type, DataTypes types = {result_type}; cctz::time_zone timezone_obj; TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, timezone_obj); - RETURN_IF_ERROR(convert_from_arrow_batch(result, types, &result_block, timezone_obj)); + RETURN_IF_ERROR(PythonArrowBlockConvertor(result->schema(), timezone_obj) + .convert_from_arrow(result, types, &result_block)); // Insert the result value into output column if (result_block.rows() != 1) { diff --git a/be/src/exprs/function/function_python_udf.cpp b/be/src/exprs/function/function_python_udf.cpp index 6c8c292819aa9c..0660c379bba778 100644 --- a/be/src/exprs/function/function_python_udf.cpp +++ b/be/src/exprs/function/function_python_udf.cpp @@ -135,17 +135,23 @@ Status PythonFunctionCall::execute_impl(FunctionContext* context, Block& block, input_block.insert(block.get_by_position(arguments[i])); } - std::shared_ptr schema; - RETURN_IF_ERROR( - get_arrow_schema_from_block(input_block, &schema, TimezoneUtils::default_time_zone)); std::shared_ptr input_batch; std::shared_ptr output_batch; - cctz::time_zone _timezone_obj; // default UTC + cctz::time_zone timezone_obj; + // Python schemas use the Doris default zone, so the writer must use the same zone; otherwise + // DATETIMEV2 metadata and encoded values describe different instants. + if (!TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, timezone_obj)) { + return Status::InternalError("Failed to resolve the default Python UDF timezone"); + } + PythonArrowBlockConvertor converter(input_block, TimezoneUtils::default_time_zone, + timezone_obj); + RETURN_IF_ERROR(converter.init()); if (arguments.empty()) { - RETURN_IF_ERROR(make_zero_column_arrow_batch(schema, input_rows, &input_batch)); + RETURN_IF_ERROR( + make_zero_column_arrow_batch(converter.arrow_schema(), input_rows, &input_batch)); } else { - RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema, arrow::default_memory_pool(), - &input_batch, _timezone_obj)); + RETURN_IF_ERROR(converter.convert_to_arrow(input_block, arrow::default_memory_pool(), + &input_batch)); } RETURN_IF_ERROR(client->evaluate(*input_batch, &output_batch)); int64_t output_rows = output_batch->num_rows(); @@ -160,8 +166,8 @@ Status PythonFunctionCall::execute_impl(FunctionContext* context, Block& block, "Python UDF output rows {} not equal to input rows {}", output_rows, input_rows)); } - RETURN_IF_ERROR( - convert_from_arrow_batch(output_batch, {_return_type}, &output_block, _timezone_obj)); + RETURN_IF_ERROR(PythonArrowBlockConvertor(output_batch->schema(), timezone_obj) + .convert_from_arrow(output_batch, {_return_type}, &output_block)); DCHECK_EQ(output_block.columns(), 1); block.replace_by_position(result, std::move(output_block.get_by_position(0).column)); return Status::OK(); diff --git a/be/src/exprs/table_function/python_udtf_function.cpp b/be/src/exprs/table_function/python_udtf_function.cpp index 9f917ab94fafd3..5352730a1ecf54 100644 --- a/be/src/exprs/table_function/python_udtf_function.cpp +++ b/be/src/exprs/table_function/python_udtf_function.cpp @@ -134,16 +134,16 @@ Status PythonUDTFFunction::process_init(Block* block, RuntimeState* state) { input_block.insert(block->get_by_position(child_column_idxs[i])); } int64_t input_rows = block->rows(); - std::shared_ptr input_schema; std::shared_ptr input_batch; - RETURN_IF_ERROR(get_arrow_schema_from_block(input_block, &input_schema, - TimezoneUtils::default_time_zone)); + PythonArrowBlockConvertor converter(input_block, TimezoneUtils::default_time_zone, + _timezone_obj); + RETURN_IF_ERROR(converter.init()); if (child_column_idxs.empty()) { - RETURN_IF_ERROR(make_zero_column_arrow_batch(input_schema, input_rows, &input_batch)); + RETURN_IF_ERROR( + make_zero_column_arrow_batch(converter.arrow_schema(), input_rows, &input_batch)); } else { - RETURN_IF_ERROR(convert_to_arrow_batch(input_block, input_schema, - arrow::default_memory_pool(), &input_batch, - _timezone_obj)); + RETURN_IF_ERROR(converter.convert_to_arrow(input_block, arrow::default_memory_pool(), + &input_batch)); } // Step 3: Call Python UDTF to evaluate all rows at once (similar to Java UDTF's JNI call) diff --git a/be/src/format/arrow/arrow_block_convertor.cpp b/be/src/format/arrow/arrow_block_convertor.cpp index d7fc3349f05f8e..0046ad3dacc328 100644 --- a/be/src/format/arrow/arrow_block_convertor.cpp +++ b/be/src/format/arrow/arrow_block_convertor.cpp @@ -41,18 +41,18 @@ #include #include -#include "common/cast_set.h" #include "common/status.h" #include "core/block/column_with_type_and_name.h" #include "core/column/column.h" -#include "core/column/column_nullable.h" -#include "core/column/column_string.h" #include "core/data_type/data_type.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" #include "core/value/vdatetime_value.h" #include "format/arrow/arrow_row_batch.h" #include "format/arrow/arrow_utils.h" +#include "util/timezone_utils.h" namespace arrow { class Array; @@ -63,15 +63,17 @@ namespace doris { namespace { -constexpr const char* ICEBERG_ORIGINAL_TYPE_KEY = "originalType"; -constexpr const char* ICEBERG_UUID_TYPE_VALUE = "uuid"; - -bool is_iceberg_uuid_field(const std::shared_ptr& field) { - if (field == nullptr || !field->HasMetadata()) { - return false; +int hex_value(char c) { + if (c >= '0' && c <= '9') { + return c - '0'; + } + if (c >= 'a' && c <= 'f') { + return c - 'a' + 10; + } + if (c >= 'A' && c <= 'F') { + return c - 'A' + 10; } - const auto result = field->metadata()->Get(ICEBERG_ORIGINAL_TYPE_KEY); - return result.ok() && result.ValueUnsafe() == ICEBERG_UUID_TYPE_VALUE; + return -1; } bool contains_extension_type(const std::shared_ptr& type) { @@ -118,6 +120,98 @@ std::shared_ptr extension_storage_type( } } +// `type` is the Doris logical type; `plain_arrow_type` is its ordinary SerDe mapping; +// `target_type` is the Arrow representation requested by the consumer. For example, +// DATETIMEV2(6) may bind timestamp(us) without a timezone, but TIMESTAMPTZ(6) must +// retain a timezone to preserve instant semantics. Nested bindings obey the same rule. +bool is_declared_plain_arrow_binding(const DataTypePtr& type, + const std::shared_ptr& plain_arrow_type, + const std::shared_ptr& target_type) { + if (plain_arrow_type->Equals(target_type)) { + return true; + } + if (plain_arrow_type->id() == arrow::Type::TIMESTAMP && + target_type->id() == arrow::Type::TIMESTAMP) { + const auto& plain_timestamp = assert_cast(*plain_arrow_type); + const auto& target_timestamp = assert_cast(*target_type); + if (plain_timestamp.unit() != target_timestamp.unit()) { + return false; + } + const PrimitiveType primitive = remove_nullable(type)->get_primitive_type(); + // A timezone-free Arrow timestamp is a wall-clock value and is therefore only compatible + // with DATETIMEV2; TIMESTAMPTZ must always retain its instant semantics. + if (target_timestamp.timezone().empty()) { + return primitive == TYPE_DATETIMEV2; + } + cctz::time_zone target_timezone; + return TimezoneUtils::find_cctz_time_zone(target_timestamp.timezone(), target_timezone) && + target_timezone.name() == plain_timestamp.timezone(); + } + const PrimitiveType primitive = remove_nullable(type)->get_primitive_type(); + if (primitive == TYPE_ARRAY && plain_arrow_type->id() == arrow::Type::LIST && + target_type->id() == arrow::Type::LIST) { + const auto& array = assert_cast(*remove_nullable(type)); + const auto& plain_list = assert_cast(*plain_arrow_type); + const auto& target_list = assert_cast(*target_type); + return plain_list.value_field() + ->WithType(target_list.value_type()) + ->Equals(target_list.value_field()) && + is_declared_plain_arrow_binding(array.get_nested_type(), plain_list.value_type(), + target_list.value_type()); + } + if (primitive == TYPE_MAP && plain_arrow_type->id() == arrow::Type::MAP && + target_type->id() == arrow::Type::MAP) { + const auto& map = assert_cast(*remove_nullable(type)); + const auto& plain_map = assert_cast(*plain_arrow_type); + const auto& target_map = assert_cast(*target_type); + return plain_map.keys_sorted() == target_map.keys_sorted() && + plain_map.key_field() + ->WithType(target_map.key_type()) + ->Equals(target_map.key_field()) && + plain_map.item_field() + ->WithType(target_map.item_type()) + ->Equals(target_map.item_field()) && + is_declared_plain_arrow_binding(map.get_key_type(), plain_map.key_type(), + target_map.key_type()) && + is_declared_plain_arrow_binding(map.get_value_type(), plain_map.item_type(), + target_map.item_type()); + } + if (primitive == TYPE_STRUCT && plain_arrow_type->id() == arrow::Type::STRUCT && + target_type->id() == arrow::Type::STRUCT) { + const auto& structure = assert_cast(*remove_nullable(type)); + if (plain_arrow_type->num_fields() != target_type->num_fields() || + structure.get_elements().size() != static_cast(target_type->num_fields())) { + return false; + } + for (int i = 0; i < target_type->num_fields(); ++i) { + const auto& plain_field = plain_arrow_type->field(i); + const auto& target_field = target_type->field(i); + if (!plain_field->WithType(target_field->type())->Equals(target_field) || + !is_declared_plain_arrow_binding(structure.get_element(i), plain_field->type(), + target_field->type())) { + return false; + } + } + return true; + } + if (is_string_type(primitive)) { + return target_type->id() == arrow::Type::STRING || + target_type->id() == arrow::Type::LARGE_STRING || + target_type->id() == arrow::Type::BINARY || + target_type->id() == arrow::Type::LARGE_BINARY; + } + if (primitive == TYPE_VARBINARY) { + return target_type->id() == arrow::Type::STRING || + target_type->id() == arrow::Type::BINARY || + target_type->id() == arrow::Type::LARGE_BINARY; + } + if (primitive == TYPE_VARIANT) { + return target_type->id() == arrow::Type::STRING || + target_type->id() == arrow::Type::LARGE_STRING; + } + return false; +} + Status wrap_extension_arrays(const std::shared_ptr& target_type, const std::shared_ptr& storage_array, std::shared_ptr* result) { @@ -166,19 +260,6 @@ Status wrap_extension_arrays(const std::shared_ptr& target_type return Status::OK(); } -int hex_value(char c) { - if (c >= '0' && c <= '9') { - return c - '0'; - } - if (c >= 'a' && c <= 'f') { - return c - 'a' + 10; - } - if (c >= 'A' && c <= 'F') { - return c - 'A' + 10; - } - return -1; -} - } // namespace Status parse_iceberg_uuid_to_bytes(StringRef uuid, std::array* bytes) { @@ -223,79 +304,82 @@ Status parse_iceberg_uuid_to_bytes(StringRef uuid, std::array* byte return Status::OK(); } -namespace { - -Status write_iceberg_uuid_string_column_to_arrow(const IColumn& column, const DataTypePtr& type, - arrow::ArrayBuilder* array_builder, int64_t start, - int64_t end) { - if (array_builder->type()->id() != arrow::Type::FIXED_SIZE_BINARY) { - return Status::InvalidArgument("Iceberg UUID must be written to fixed size binary"); - } - const int byte_width = - static_cast(*array_builder->type()).byte_width(); - if (byte_width != 16) { - return Status::InvalidArgument("Iceberg UUID expects 16 bytes, got {}", byte_width); - } - - auto& builder = assert_cast(*array_builder); - const IColumn* data_column = &column; - const NullMap* null_map = nullptr; - if (type->is_nullable()) { - const auto& nullable_column = assert_cast(column); - data_column = &nullable_column.get_nested_column(); - null_map = &nullable_column.get_null_map_data(); - } - if (!data_column->is_column_string()) { +Status ArrowBlockConvertor::write_plain_arrow_column(const std::shared_ptr& type, + const DataTypeSerDe& serde, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, + int64_t start, int64_t end, + const cctz::time_zone& ctz) const { + std::shared_ptr plain_arrow_type; + RETURN_IF_ERROR(convert_to_arrow_type(type, &plain_arrow_type, ctz.name())); + const auto storage_type = extension_storage_type(field->type()); + // This is an exact binding check selected by the target converter, not a recovery path. A + // mismatch returns without invoking SerDe, and a SerDe error is never retried elsewhere. + if (!is_declared_plain_arrow_binding(type, plain_arrow_type, storage_type)) { return Status::InvalidArgument( - "Iceberg UUID string conversion expects string column, got {}", - data_column->get_name()); + "Plain Arrow writer is not bound for Doris type {} and Arrow field {}", + type->get_name(), field->ToString()); } + return serde.write_column_to_arrow(column, null_map, array_builder, start, end, ctz); +} - const auto& string_column = assert_cast(*data_column); - const auto begin_row = cast_set(start); - const auto end_row = cast_set(end); - for (size_t row = begin_row; row < end_row; ++row) { - if (null_map != nullptr && (*null_map)[row]) { - RETURN_IF_ERROR(checkArrowStatus(builder.AppendNull(), column, builder)); - continue; - } - std::array bytes; - RETURN_IF_ERROR(parse_iceberg_uuid_to_bytes(string_column.get_data_at(row), &bytes)); - RETURN_IF_ERROR(checkArrowStatus(builder.Append(bytes.data()), column, builder)); +Status DorisArrowBlockConvertor::write_column(const std::shared_ptr& type, + const DataTypeSerDe& serde, const IColumn& column, + const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const { + return write_plain_arrow_column(type, serde, column, null_map, field, array_builder, start, end, + ctz); +} + +Status ArrowBlockConvertor::init() { + if (_arrow_schema == nullptr) { + return Status::InvalidArgument("Arrow converter schema is not initialized"); } return Status::OK(); } -} // namespace - -Status FromBlockToRecordBatchConverter::convert(std::shared_ptr* out) { - int num_fields = _schema->num_fields(); - if (_block.columns() != num_fields) { - return Status::InvalidArgument("number fields not match"); +Status DorisArrowBlockConvertor::init() { + if (_arrow_schema == nullptr) { + // cctz names fixed offsets as "Fixed/UTC+HH:MM:SS", which is not the Arrow + // protocol label. Keep the declared name so Python metadata and batches agree. + RETURN_IF_ERROR(get_arrow_schema_from_block(_header, &_arrow_schema, _timezone_name, + _datetime_naive)); } + return ArrowBlockConvertor::init(); +} - // Calculate actual row range to convert - size_t actual_start = _row_range_start; - size_t actual_rows = _row_range_end > 0 ? (_row_range_end - _row_range_start) - : (_block.rows() - _row_range_start); +Status ArrowBlockConvertor::convert_from_arrow(const std::shared_ptr& batch, + const DataTypes& types, Block* block) const { + // Iceberg and Paimon physical layouts are not the generic Arrow SerDe read contract. + return Status::NotSupported("This Arrow converter does not support reading"); +} - // Validate range - if (actual_start + actual_rows > _block.rows()) { - return Status::InvalidArgument( - "Row range out of bounds: start={}, num_rows={}, block_rows={}", actual_start, - actual_rows, _block.rows()); +Status ArrowBlockConvertor::convert_to_arrow(const Block& block, arrow::MemoryPool* pool, + std::shared_ptr* out, + size_t start_row, size_t end_row) const { + if (_arrow_schema == nullptr) { + return Status::InvalidArgument("Arrow converter schema is not initialized"); } - - _arrays.resize(num_fields); - + const auto& schema = _arrow_schema; + int num_fields = schema->num_fields(); + if (block.columns() != num_fields) { + return Status::InvalidArgument("number fields not match"); + } + const size_t actual_end = end_row == 0 ? block.rows() : end_row; + // Validate endpoints before subtraction: an inverted unsigned range otherwise wraps. + if (start_row > actual_end || actual_end > block.rows()) { + return Status::InvalidArgument("Row range out of bounds: start={}, end={}, block_rows={}", + start_row, actual_end, block.rows()); + } + const size_t actual_rows = actual_end - start_row; + std::vector> arrays(num_fields); for (int idx = 0; idx < num_fields; ++idx) { - _cur_field_idx = idx; - _cur_start = actual_start; - _cur_rows = actual_rows; - _cur_col = _block.get_by_position(idx).column; - _cur_type = _block.get_by_position(idx).type; - auto column = _cur_col->convert_to_full_column_if_const(); - auto target_arrow_type = _schema->field(idx)->type(); + const auto& entry = block.get_by_position(idx); + auto column = entry.column->convert_to_full_column_if_const(); + auto target_arrow_type = schema->field(idx)->type(); const bool has_extension = contains_extension_type(target_arrow_type); auto builder_arrow_type = has_extension ? extension_storage_type(target_arrow_type) : target_arrow_type; @@ -307,82 +391,58 @@ Status FromBlockToRecordBatchConverter::convert(std::shared_ptr builder; - auto arrow_st = arrow::MakeBuilder(_pool, builder_arrow_type, &builder); + auto arrow_st = arrow::MakeBuilder(pool, builder_arrow_type, &builder); if (!arrow_st.ok()) { return to_doris_status(arrow_st); } - _cur_builder = builder.get(); try { - if (is_iceberg_uuid_field(_schema->field(idx)) && - is_string_type(remove_nullable(_cur_type)->get_primitive_type())) { - RETURN_IF_ERROR(write_iceberg_uuid_string_column_to_arrow( - *column, _cur_type, _cur_builder, _cur_start, _cur_start + _cur_rows)); - } else { - RETURN_IF_ERROR(_cur_type->get_serde()->write_column_to_arrow( - *column, nullptr, _cur_builder, _cur_start, _cur_start + _cur_rows, - _timezone_obj)); - } + const auto serde = entry.type->get_serde(); + RETURN_IF_ERROR(write_column(entry.type, *serde, *column, nullptr, schema->field(idx), + builder.get(), start_row, actual_end, _timezone)); } catch (std::exception& e) { return Status::InternalError( "Fail to convert block data to arrow data, type: {}, name: {}, error: {}", - _cur_type->get_name(), _block.get_by_position(idx).name, e.what()); + entry.type->get_name(), entry.name, e.what()); } std::shared_ptr storage_array; - arrow_st = _cur_builder->Finish(&storage_array); + arrow_st = builder->Finish(&storage_array); if (!arrow_st.ok()) { return to_doris_status(arrow_st); } if (has_extension) { - RETURN_IF_ERROR(wrap_extension_arrays(target_arrow_type, storage_array, - &_arrays[_cur_field_idx])); + RETURN_IF_ERROR(wrap_extension_arrays(target_arrow_type, storage_array, &arrays[idx])); } else { - _arrays[_cur_field_idx] = std::move(storage_array); + arrays[idx] = std::move(storage_array); } } - *out = arrow::RecordBatch::Make(_schema, actual_rows, std::move(_arrays)); + *out = arrow::RecordBatch::Make(schema, actual_rows, std::move(arrays)); return Status::OK(); } -Status FromRecordBatchToBlockConverter::convert(Block* block) { +Status DorisArrowBlockConvertor::convert_from_arrow( + const std::shared_ptr& batch, const DataTypes& types, + Block* block) const { DCHECK(block); - int num_fields = _batch->num_columns(); - if ((size_t)num_fields != _types.size()) { + int num_fields = batch->num_columns(); + if ((size_t)num_fields != types.size()) { return Status::InvalidArgument("number fields not match"); } - - int64_t num_rows = _batch->num_rows(); - _columns.reserve(num_fields); - + int64_t num_rows = batch->num_rows(); + ColumnsWithTypeAndName columns; + columns.reserve(num_fields); for (int idx = 0; idx < num_fields; ++idx) { - auto doris_type = _types[idx]; + auto doris_type = types[idx]; auto doris_column = doris_type->create_column(); - auto arrow_column = _batch->column(idx); + auto arrow_column = batch->column(idx); DCHECK_EQ(arrow_column->length(), num_rows); RETURN_IF_ERROR(doris_type->get_serde()->read_column_from_arrow( - *doris_column, &*arrow_column, 0, num_rows, _timezone_obj)); - _columns.emplace_back(std::move(doris_column), std::move(doris_type), std::to_string(idx)); + *doris_column, &*arrow_column, 0, num_rows, _timezone)); + columns.emplace_back(std::move(doris_column), std::move(doris_type), std::to_string(idx)); } - - block->swap(_columns); + block->swap(columns); return Status::OK(); } -Status convert_to_arrow_batch(const Block& block, const std::shared_ptr& schema, - arrow::MemoryPool* pool, std::shared_ptr* result, - const cctz::time_zone& timezone_obj) { - FromBlockToRecordBatchConverter converter(block, schema, pool, timezone_obj); - return converter.convert(result); -} - -Status convert_to_arrow_batch(const Block& block, const std::shared_ptr& schema, - arrow::MemoryPool* pool, std::shared_ptr* result, - const cctz::time_zone& timezone_obj, size_t start_row, - size_t end_row) { - FromBlockToRecordBatchConverter converter(block, schema, pool, timezone_obj, start_row, - end_row); - return converter.convert(result); -} - Status make_zero_column_arrow_batch(const std::shared_ptr& schema, int64_t rows, std::shared_ptr* result) { if (schema->num_fields() != 0) { @@ -392,12 +452,5 @@ Status make_zero_column_arrow_batch(const std::shared_ptr& schema return Status::OK(); } -Status convert_from_arrow_batch(const std::shared_ptr& batch, - const DataTypes& types, Block* block, - const cctz::time_zone& timezone_obj) { - FromRecordBatchToBlockConverter converter(batch, types, timezone_obj); - return converter.convert(block); -} - #include "common/compile_check_end.h" } // namespace doris diff --git a/be/src/format/arrow/arrow_block_convertor.h b/be/src/format/arrow/arrow_block_convertor.h index 2ba3df94cd948c..ec1083505a9c16 100644 --- a/be/src/format/arrow/arrow_block_convertor.h +++ b/be/src/format/arrow/arrow_block_convertor.h @@ -43,89 +43,86 @@ class Schema; namespace doris { +// ORC and Arrow Iceberg writers share this parser so textual and binary UUID inputs always use +// the same canonical 16-byte representation. Status parse_iceberg_uuid_to_bytes(StringRef uuid, std::array* bytes); -class FromBlockToRecordBatchConverter { +// One converter owns both batch orchestration and its format-specific column bindings. +// Schema and timezone belong to the instance so one writer cannot borrow another protocol's bindings. +class ArrowBlockConvertor { public: - FromBlockToRecordBatchConverter(const Block& block, - const std::shared_ptr& schema, - arrow::MemoryPool* pool, const cctz::time_zone& timezone_obj) - : _block(block), - _schema(schema), - _pool(pool), - _cur_field_idx(-1), - _timezone_obj(timezone_obj), - _row_range_start(0), - _row_range_end(0) {} - - FromBlockToRecordBatchConverter(const Block& block, - const std::shared_ptr& schema, - arrow::MemoryPool* pool, const cctz::time_zone& timezone_obj, - size_t start_row, size_t end_row) - : _block(block), - _schema(schema), - _pool(pool), - _cur_field_idx(-1), - _timezone_obj(timezone_obj), - _row_range_start(start_row), - _row_range_end(end_row) {} - - ~FromBlockToRecordBatchConverter() = default; - - Status convert(std::shared_ptr* out); - -private: - const Block& _block; - const std::shared_ptr& _schema; - arrow::MemoryPool* _pool; - - size_t _cur_field_idx; - size_t _cur_start; - size_t _cur_rows; - ColumnPtr _cur_col; - DataTypePtr _cur_type; - arrow::ArrayBuilder* _cur_builder = nullptr; - - const cctz::time_zone& _timezone_obj; - - // Row range for zero-copy slicing (0 means use all rows from _row_range_start) - size_t _row_range_start; - size_t _row_range_end; - - std::vector> _arrays; + ArrowBlockConvertor(std::shared_ptr schema, const cctz::time_zone& timezone) + : _arrow_schema(std::move(schema)), _timezone(timezone) {} + virtual ~ArrowBlockConvertor() = default; + + virtual Status init(); + const std::shared_ptr& arrow_schema() const { return _arrow_schema; } + + Status convert_to_arrow(const Block& block, arrow::MemoryPool* pool, + std::shared_ptr* result, size_t start_row = 0, + size_t end_row = 0) const; + + virtual Status convert_from_arrow(const std::shared_ptr& batch, + const DataTypes& types, Block* block) const; + +protected: + std::shared_ptr _arrow_schema; + const cctz::time_zone _timezone; + + virtual Status write_column(const std::shared_ptr& type, + const DataTypeSerDe& serde, const IColumn& column, + const NullMap* null_map, const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const = 0; + + Status write_plain_arrow_column(const std::shared_ptr& type, + const DataTypeSerDe& serde, const IColumn& column, + const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const; }; -class FromRecordBatchToBlockConverter { +// The ordinary Doris Arrow protocol is shared explicitly by its consumers, never selected +// as a fallback for table formats with different timestamp or nested-type semantics. +class DorisArrowBlockConvertor : public ArrowBlockConvertor { public: - FromRecordBatchToBlockConverter(const std::shared_ptr& batch, - const DataTypes& types, const cctz::time_zone& timezone_obj) - : _batch(batch), _types(types), _timezone_obj(timezone_obj) {} - - ~FromRecordBatchToBlockConverter() = default; - - Status convert(Block* block); + using ArrowBlockConvertor::ArrowBlockConvertor; + DorisArrowBlockConvertor(const Block& header, std::string timezone_name, + const cctz::time_zone& timezone, bool datetime_naive = false) + : ArrowBlockConvertor(nullptr, timezone), + _header(header.clone_empty()), + _timezone_name(std::move(timezone_name)), + _datetime_naive(datetime_naive) {} + + Status init() override; + Status convert_from_arrow(const std::shared_ptr& batch, + const DataTypes& types, Block* block) const override; + +protected: + Status write_column(const std::shared_ptr& type, const DataTypeSerDe& serde, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const override; private: - const std::shared_ptr& _batch; - const DataTypes& _types; - const cctz::time_zone& _timezone_obj; - ColumnsWithTypeAndName _columns; + Block _header; + std::string _timezone_name; + bool _datetime_naive = false; }; -Status convert_to_arrow_batch(const Block& block, const std::shared_ptr& schema, - arrow::MemoryPool* pool, std::shared_ptr* result, - const cctz::time_zone& timezone_obj); +class ArrowFlightArrowBlockConvertor final : public DorisArrowBlockConvertor { +public: + using DorisArrowBlockConvertor::DorisArrowBlockConvertor; +}; -Status convert_to_arrow_batch(const Block& block, const std::shared_ptr& schema, - arrow::MemoryPool* pool, std::shared_ptr* result, - const cctz::time_zone& timezone_obj, size_t start_row, - size_t end_row); +class PythonArrowBlockConvertor final : public DorisArrowBlockConvertor { +public: + using DorisArrowBlockConvertor::DorisArrowBlockConvertor; +}; Status make_zero_column_arrow_batch(const std::shared_ptr& schema, int64_t rows, std::shared_ptr* result); -Status convert_from_arrow_batch(const std::shared_ptr& batch, - const DataTypes& types, Block* block, - const cctz::time_zone& timezone_obj); - } // namespace doris diff --git a/be/src/format/arrow/arrow_row_batch.cpp b/be/src/format/arrow/arrow_row_batch.cpp index cc88bc18b92f29..f291cf33ac38db 100644 --- a/be/src/format/arrow/arrow_row_batch.cpp +++ b/be/src/format/arrow/arrow_row_batch.cpp @@ -114,7 +114,10 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type, if (type->get_primitive_type() == TYPE_DATETIMEV2 && datetime_naive) { *result = std::make_shared(time_unit); } else { - *result = std::make_shared(time_unit, timezone); + // Arrow clients resolve timezone metadata as an IANA name; use the canonical UTC + // name instead of the ISO-8601 "Z" alias without changing the encoded instant. + *result = std::make_shared(time_unit, + timezone == "Z" ? "UTC" : timezone); } break; } diff --git a/be/src/format/parquet/parquet_arrow_block_convertor.cpp b/be/src/format/parquet/parquet_arrow_block_convertor.cpp new file mode 100644 index 00000000000000..6dc67c1aa432e6 --- /dev/null +++ b/be/src/format/parquet/parquet_arrow_block_convertor.cpp @@ -0,0 +1,55 @@ +// 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/parquet/parquet_arrow_block_convertor.h" + +#include + +#include "format/arrow/arrow_row_batch.h" + +namespace doris { +#include "common/compile_check_begin.h" + +Status ParquetArrowBlockConvertor::init() { + if (_types.size() != _names.size()) { + return Status::InvalidArgument("Parquet column names and types must have the same size"); + } + std::vector> fields; + fields.reserve(_types.size()); + // Retain the declared Arrow timezone label; cctz's fixed-offset name is internal. + // Preserve the existing timezone-bearing DATETIMEV2 schema for both Parquet encodings. + for (size_t i = 0; i < _types.size(); ++i) { + std::shared_ptr type; + RETURN_IF_ERROR(convert_to_arrow_type(_types[i], &type, _timezone_name)); + fields.emplace_back(arrow::field(_names[i], type, _types[i]->is_nullable())); + } + _arrow_schema = arrow::schema(std::move(fields)); + return Status::OK(); +} + +Status ParquetArrowBlockConvertor::write_column(const DataTypePtr& type, const DataTypeSerDe& serde, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* builder, int64_t start, + int64_t end, + const cctz::time_zone& timezone) const { + return write_plain_arrow_column(type, serde, column, null_map, field, builder, start, end, + timezone); +} + +#include "common/compile_check_end.h" +} // namespace doris diff --git a/be/src/format/parquet/parquet_arrow_block_convertor.h b/be/src/format/parquet/parquet_arrow_block_convertor.h new file mode 100644 index 00000000000000..8889a38c498d09 --- /dev/null +++ b/be/src/format/parquet/parquet_arrow_block_convertor.h @@ -0,0 +1,46 @@ +// 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. + +#pragma once + +#include "format/arrow/arrow_block_convertor.h" + +namespace doris { + +class ParquetArrowBlockConvertor : public ArrowBlockConvertor { +public: + ParquetArrowBlockConvertor(DataTypes types, std::vector names, + std::string timezone_name, const cctz::time_zone& timezone) + : ArrowBlockConvertor(nullptr, timezone), + _types(std::move(types)), + _names(std::move(names)), + _timezone_name(std::move(timezone_name)) {} + Status init() override; + +protected: + Status write_column(const DataTypePtr& type, const DataTypeSerDe& serde, const IColumn& column, + const NullMap* null_map, const std::shared_ptr& field, + arrow::ArrayBuilder* builder, int64_t start, int64_t end, + const cctz::time_zone& timezone) const override; + +private: + DataTypes _types; + std::vector _names; + std::string _timezone_name; +}; + +} // namespace doris diff --git a/be/src/format/table/hive/hive_arrow_block_convertor.h b/be/src/format/table/hive/hive_arrow_block_convertor.h new file mode 100644 index 00000000000000..490bcf9a685fb1 --- /dev/null +++ b/be/src/format/table/hive/hive_arrow_block_convertor.h @@ -0,0 +1,31 @@ +// 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. + +#pragma once + +#include "format/parquet/parquet_arrow_block_convertor.h" + +namespace doris::hive { + +// Hive currently shares Parquet's physical bindings, with the writer timezone supplied +// explicitly by the writer. Keep its protocol separate from Arrow Flight and other tables. +class HiveArrowBlockConvertor final : public ParquetArrowBlockConvertor { +public: + using ParquetArrowBlockConvertor::ParquetArrowBlockConvertor; +}; + +} // namespace doris::hive diff --git a/be/src/format/table/iceberg/iceberg_arrow_block_convertor.cpp b/be/src/format/table/iceberg/iceberg_arrow_block_convertor.cpp new file mode 100644 index 00000000000000..f5260b76e6d22b --- /dev/null +++ b/be/src/format/table/iceberg/iceberg_arrow_block_convertor.cpp @@ -0,0 +1,57 @@ +// 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/table/iceberg/iceberg_arrow_block_convertor.h" + +#include +#include +#include + +#include "format/table/iceberg/arrow_schema_util.h" + +namespace doris::iceberg { +#include "common/compile_check_begin.h" + +Status IcebergArrowBlockConvertor::init() { + if (_schema == nullptr) { + return ArrowBlockConvertor::init(); + } + // Field IDs, Variant storage and timestamp bindings must share the same target schema. + std::vector> fields; + // Arrow consumers need the declared label, not cctz's internal fixed-offset name. + RETURN_IF_ERROR(ArrowSchemaUtil::convert(_schema, _timezone_name, fields)); + _arrow_schema = arrow::schema(std::move(fields)); + if (!_schema_json.empty()) { + _arrow_schema = _arrow_schema->WithMetadata( + arrow::KeyValueMetadata::Make({"iceberg.schema"}, {_schema_json})); + } + return Status::OK(); +} + +Status IcebergArrowBlockConvertor::write_column(const std::shared_ptr& type, + const DataTypeSerDe& serde, const IColumn& column, + const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const { + return serde.write_column_to_iceberg_arrow(type, column, null_map, + field->WithType(array_builder->type()), + array_builder, start, end, ctz); +} + +#include "common/compile_check_end.h" +} // namespace doris::iceberg diff --git a/be/src/format/table/iceberg/iceberg_arrow_block_convertor.h b/be/src/format/table/iceberg/iceberg_arrow_block_convertor.h new file mode 100644 index 00000000000000..1d2c8d93f4c074 --- /dev/null +++ b/be/src/format/table/iceberg/iceberg_arrow_block_convertor.h @@ -0,0 +1,51 @@ +// 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. + +#pragma once + +#include "format/arrow/arrow_block_convertor.h" + +namespace doris::iceberg { + +class Schema; + +class IcebergArrowBlockConvertor final : public ArrowBlockConvertor { +public: + using ArrowBlockConvertor::ArrowBlockConvertor; + IcebergArrowBlockConvertor(const Schema& schema, const std::string* schema_json, + std::string timezone_name, const cctz::time_zone& timezone) + : ArrowBlockConvertor(nullptr, timezone), + _schema(&schema), + _schema_json(schema_json == nullptr ? "" : *schema_json), + _timezone_name(std::move(timezone_name)) {} + + Status init() override; + +protected: + Status write_column(const std::shared_ptr& type, const DataTypeSerDe& serde, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const override; + +private: + const Schema* _schema = nullptr; + std::string _schema_json; + std::string _timezone_name; +}; + +} // namespace doris::iceberg diff --git a/be/src/format/table/paimon/paimon_arrow_block_convertor.cpp b/be/src/format/table/paimon/paimon_arrow_block_convertor.cpp new file mode 100644 index 00000000000000..9404d622938492 --- /dev/null +++ b/be/src/format/table/paimon/paimon_arrow_block_convertor.cpp @@ -0,0 +1,58 @@ +// 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/table/paimon/paimon_arrow_block_convertor.h" + +#include +#include +#include +#include + +namespace doris::paimon { +#include "common/compile_check_begin.h" + +Status PaimonArrowBlockConvertor::init() { + if (_arrow_schema != nullptr) { + return Status::OK(); + } + // Decode the pinned table schema here; rebuilding it from Doris types would lose + // nested nullability, timestamp precision and Paimon's physical Variant layout. + auto input = std::make_shared( + arrow::Buffer::FromString(_serialized_schema)); + auto reader = arrow::ipc::RecordBatchStreamReader::Open(input); + if (!reader.ok()) { + return Status::InvalidArgument("Failed to deserialize Paimon Arrow schema: {}", + reader.status().ToString()); + } + _arrow_schema = reader.ValueOrDie()->schema(); + _serialized_schema.clear(); + return Status::OK(); +} + +Status PaimonArrowBlockConvertor::write_column(const std::shared_ptr& type, + const DataTypeSerDe& serde, const IColumn& column, + const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, + int64_t end, const cctz::time_zone& ctz) const { + return serde.write_column_to_paimon_arrow(type, column, null_map, + field->WithType(array_builder->type()), array_builder, + start, end, ctz); +} + +#include "common/compile_check_end.h" +} // namespace doris::paimon diff --git a/be/src/format/table/paimon/paimon_arrow_block_convertor.h b/be/src/format/table/paimon/paimon_arrow_block_convertor.h new file mode 100644 index 00000000000000..caa6501e0a65e7 --- /dev/null +++ b/be/src/format/table/paimon/paimon_arrow_block_convertor.h @@ -0,0 +1,47 @@ +// 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. + +#pragma once + +#include "format/arrow/arrow_block_convertor.h" + +namespace doris::paimon { + +// Paimon reads use native Parquet/ORC or the JNI scanner, which apply the table's +// timestamp and nested-type semantics. This adapter implements the Arrow write protocol only; +// convert_from_arrow deliberately inherits NotSupported instead of a generic SerDe fallback. +class PaimonArrowBlockConvertor final : public ArrowBlockConvertor { +public: + using ArrowBlockConvertor::ArrowBlockConvertor; + PaimonArrowBlockConvertor(std::string serialized_schema, const cctz::time_zone& timezone) + : ArrowBlockConvertor(nullptr, timezone), + _serialized_schema(std::move(serialized_schema)) {} + + Status init() override; + +protected: + Status write_column(const std::shared_ptr& type, const DataTypeSerDe& serde, + const IColumn& column, const NullMap* null_map, + const std::shared_ptr& field, + arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, + const cctz::time_zone& ctz) const override; + +private: + std::string _serialized_schema; +}; + +} // namespace doris::paimon diff --git a/be/src/format/transformer/vfile_format_transformer_factory.cpp b/be/src/format/transformer/vfile_format_transformer_factory.cpp index 8d714c21387041..501dadc0c2ba6a 100644 --- a/be/src/format/transformer/vfile_format_transformer_factory.cpp +++ b/be/src/format/transformer/vfile_format_transformer_factory.cpp @@ -25,7 +25,7 @@ #include "format/transformer/vcsv_transformer.h" #include "format/transformer/vjni_format_transformer.h" #include "format/transformer/vorc_transformer.h" -#include "format/transformer/vparquet_transformer.h" +#include "format/transformer/vparquet_writer.h" namespace doris { @@ -79,7 +79,7 @@ Status create_tvf_format_transformer(const TTVFTableSink& tvf_sink, RuntimeState parquet_schemas.push_back(schema); } } - result->reset(new VParquetTransformer( + result->reset(new VParquetWriter( state, file_writer, output_vexpr_ctxs, parquet_schemas, false, {TParquetCompressionType::SNAPPY, TParquetVersion::PARQUET_1_0, false, false})); break; diff --git a/be/src/format/transformer/vhive_parquet_writer.h b/be/src/format/transformer/vhive_parquet_writer.h new file mode 100644 index 00000000000000..a591f44486387b --- /dev/null +++ b/be/src/format/transformer/vhive_parquet_writer.h @@ -0,0 +1,38 @@ +// 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. + +#pragma once + +#include "format/table/hive/hive_arrow_block_convertor.h" +#include "format/transformer/vparquet_writer.h" + +namespace doris { + +class VHiveParquetWriter final : public VParquetWriter { +public: + using VParquetWriter::VParquetWriter; + +protected: + std::unique_ptr _create_arrow_block_convertor( + DataTypes types, std::vector names, const std::string& timezone_name, + const cctz::time_zone& timezone) const override { + return std::make_unique(std::move(types), std::move(names), + timezone_name, timezone); + } +}; + +} // namespace doris diff --git a/be/src/format/transformer/viceberg_parquet_writer.cpp b/be/src/format/transformer/viceberg_parquet_writer.cpp new file mode 100644 index 00000000000000..df62430cbb263c --- /dev/null +++ b/be/src/format/transformer/viceberg_parquet_writer.cpp @@ -0,0 +1,133 @@ +// 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/transformer/viceberg_parquet_writer.h" + +#include +#include + +#include + +#include "format/table/iceberg/iceberg_arrow_block_convertor.h" +#include "format/table/parquet_utils.h" +#include "runtime/runtime_state.h" + +namespace doris { +#include "common/compile_check_begin.h" + +VIcebergParquetWriter::VIcebergParquetWriter(RuntimeState* state, io::FileWriter* file_writer, + const VExprContextSPtrs& output_vexpr_ctxs, + std::vector column_names, + bool output_object_data, + const ParquetFileOptions& parquet_options, + const std::string* iceberg_schema_json, + const iceberg::Schema& iceberg_schema) + : VParquetWriter(state, file_writer, output_vexpr_ctxs, std::move(column_names), + output_object_data, parquet_options), + _iceberg_schema(iceberg_schema), + _iceberg_schema_json(iceberg_schema_json == nullptr ? "" : *iceberg_schema_json) {} + +std::unique_ptr VIcebergParquetWriter::_create_arrow_block_convertor( + DataTypes types, std::vector names, const std::string& timezone_name, + const cctz::time_zone& timezone) const { + return std::make_unique( + _iceberg_schema, &_iceberg_schema_json, timezone_name, timezone); +} + +Status VIcebergParquetWriter::collect_file_statistics_after_close(TIcebergColumnStats* stats) { + std::shared_ptr<::parquet::FileMetaData> file_metadata = _file_metadata(); + if (file_metadata == nullptr) { + return Status::InternalError("File metadata is not available"); + } + std::map column_sizes; + std::map value_counts; + std::map null_value_counts; + std::map lower_bounds; + std::map upper_bounds; + std::map> merged_column_stats; + std::unordered_set variant_field_ids; + + const int num_row_groups = file_metadata->num_row_groups(); + const int num_columns = file_metadata->num_columns(); + for (int col_idx = 0; col_idx < num_columns; ++col_idx) { + const auto& schema_node = file_metadata->schema()->Column(col_idx)->schema_node(); + const auto* parent = schema_node->parent(); + const bool is_variant_child = parent != nullptr && parent->logical_type() != nullptr && + parent->logical_type()->is_variant(); + if (is_variant_child && schema_node->name() != "metadata") { + // The value leaf has no independent Iceberg field and its byte statistics are not + // logical Variant statistics. + continue; + } + const int field_id = is_variant_child ? parent->field_id() : schema_node->field_id(); + if (field_id < 0) { + // Parquet structural leaves (including Variant children) may intentionally omit an + // Iceberg field id. Never publish them under the synthetic -1 key. + continue; + } + if (is_variant_child) { + variant_field_ids.insert(field_id); + } + + for (int rg_idx = 0; rg_idx < num_row_groups; ++rg_idx) { + auto row_group = file_metadata->RowGroup(rg_idx); + auto column_chunk = row_group->ColumnChunk(col_idx); + if (!is_variant_child) { + column_sizes[field_id] += column_chunk->total_compressed_size(); + } + + if (column_chunk->is_stats_set()) { + auto column_stat = column_chunk->statistics(); + if (!merged_column_stats.contains(field_id)) { + merged_column_stats[field_id] = column_stat; + } else { + parquet_utils::merge_stats(merged_column_stats[field_id], column_stat); + } + } + } + } + + bool has_any_null_count = false; + bool has_any_min_max = false; + for (const auto& [field_id, column_stat] : merged_column_stats) { + value_counts[field_id] = column_stat->num_values(); + if (column_stat->HasNullCount()) { + has_any_null_count = true; + int64_t null_count = column_stat->null_count(); + null_value_counts[field_id] = null_count; + value_counts[field_id] += null_count; + } + if (!variant_field_ids.contains(field_id) && column_stat->HasMinMax()) { + has_any_min_max = true; + lower_bounds[field_id] = column_stat->EncodeMin(); + upper_bounds[field_id] = column_stat->EncodeMax(); + } + } + + stats->__set_column_sizes(column_sizes); + stats->__set_value_counts(value_counts); + if (has_any_null_count) { + stats->__set_null_value_counts(null_value_counts); + } + if (has_any_min_max) { + stats->__set_lower_bounds(lower_bounds); + stats->__set_upper_bounds(upper_bounds); + } + return Status::OK(); +} + +} // namespace doris diff --git a/be/src/format/transformer/viceberg_parquet_writer.h b/be/src/format/transformer/viceberg_parquet_writer.h new file mode 100644 index 00000000000000..365dcc7cb2e321 --- /dev/null +++ b/be/src/format/transformer/viceberg_parquet_writer.h @@ -0,0 +1,49 @@ +// 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. + +#pragma once + +#include "format/table/iceberg/iceberg_arrow_block_convertor.h" +#include "format/table/iceberg/schema.h" +#include "format/transformer/vparquet_writer.h" + +namespace doris { +#include "common/compile_check_begin.h" + +class VIcebergParquetWriter final : public VParquetWriter { +public: + VIcebergParquetWriter(RuntimeState* state, io::FileWriter* file_writer, + const VExprContextSPtrs& output_vexpr_ctxs, + std::vector column_names, bool output_object_data, + const ParquetFileOptions& parquet_options, + const std::string* iceberg_schema_json, + const iceberg::Schema& iceberg_schema); + + Status collect_file_statistics_after_close(TIcebergColumnStats* stats); + +protected: + std::unique_ptr _create_arrow_block_convertor( + DataTypes types, std::vector names, const std::string& timezone_name, + const cctz::time_zone& timezone) const override; + +private: + const iceberg::Schema& _iceberg_schema; + std::string _iceberg_schema_json; +}; + +} // namespace doris +#include "common/compile_check_end.h" diff --git a/be/src/format/transformer/vjni_format_transformer.h b/be/src/format/transformer/vjni_format_transformer.h index 4a88ec21d600ee..96588c5ce89afd 100644 --- a/be/src/format/transformer/vjni_format_transformer.h +++ b/be/src/format/transformer/vjni_format_transformer.h @@ -29,7 +29,7 @@ namespace doris { /** * VJniFormatTransformer is a VFileFormatTransformer implementation that delegates * write operations to a Java-side JniWriter via JNI. It sits alongside - * VCSVTransformer/VParquetTransformer/VOrcTransformer as a peer implementation. + * VCSVTransformer/VParquetWriter/VOrcTransformer as a peer implementation. * * The Java writer class must extend org.apache.doris.common.jni.JniWriter and * follow the same constructor signature as JniScanner: (int batchSize, Map params). diff --git a/be/src/format/transformer/vorc_transformer.h b/be/src/format/transformer/vorc_transformer.h index a863d343e5d006..9ed18a4c081b7f 100644 --- a/be/src/format/transformer/vorc_transformer.h +++ b/be/src/format/transformer/vorc_transformer.h @@ -29,7 +29,7 @@ #include "core/block/block.h" #include "core/column/column_nullable.h" #include "format/table/iceberg/schema.h" -#include "format/transformer/vparquet_transformer.h" +#include "format/transformer/vparquet_writer.h" #include "orc/Type.hh" #include "orc/Writer.hh" diff --git a/be/src/format/transformer/vparquet_transformer.cpp b/be/src/format/transformer/vparquet_writer.cpp similarity index 54% rename from be/src/format/transformer/vparquet_transformer.cpp rename to be/src/format/transformer/vparquet_writer.cpp index 0900b85b94c94c..3e28ef97107533 100644 --- a/be/src/format/transformer/vparquet_transformer.cpp +++ b/be/src/format/transformer/vparquet_writer.cpp @@ -15,12 +15,11 @@ // specific language governing permissions and limitations // under the License. -#include "format/transformer/vparquet_transformer.h" +#include "format/transformer/vparquet_writer.h" #include #include #include -#include #include #include #include @@ -33,17 +32,14 @@ #include #include #include -#include #include "common/config.h" #include "common/status.h" #include "exprs/vexpr.h" #include "exprs/vexpr_context.h" -#include "format/arrow/arrow_block_convertor.h" #include "format/arrow/arrow_row_batch.h" #include "format/arrow/arrow_utils.h" -#include "format/table/iceberg/arrow_schema_util.h" -#include "format/table/parquet_utils.h" +#include "format/parquet/parquet_arrow_block_convertor.h" #include "io/fs/file_writer.h" #include "runtime/exec_env.h" #include "runtime/runtime_state.h" @@ -180,36 +176,27 @@ void ParquetBuildHelper::build_version(::parquet::WriterProperties::Builder& bui } } -VParquetTransformer::VParquetTransformer(RuntimeState* state, doris::io::FileWriter* file_writer, - const VExprContextSPtrs& output_vexpr_ctxs, - std::vector column_names, - bool output_object_data, - const ParquetFileOptions& parquet_options, - const std::string* iceberg_schema_json, - const iceberg::Schema* iceberg_schema) +VParquetWriter::VParquetWriter(RuntimeState* state, doris::io::FileWriter* file_writer, + const VExprContextSPtrs& output_vexpr_ctxs, + std::vector column_names, bool output_object_data, + const ParquetFileOptions& parquet_options) : VFileFormatTransformer(state, output_vexpr_ctxs, output_object_data), _column_names(std::move(column_names)), - _parquet_options(parquet_options), - _iceberg_schema_json(iceberg_schema_json), - _iceberg_schema(iceberg_schema) { + _parquet_options(parquet_options) { _outstream = std::shared_ptr(new ParquetOutputStream(file_writer)); } -VParquetTransformer::VParquetTransformer(RuntimeState* state, doris::io::FileWriter* file_writer, - const VExprContextSPtrs& output_vexpr_ctxs, - std::vector parquet_schemas, - bool output_object_data, - const ParquetFileOptions& parquet_options, - const std::string* iceberg_schema_json) +VParquetWriter::VParquetWriter(RuntimeState* state, doris::io::FileWriter* file_writer, + const VExprContextSPtrs& output_vexpr_ctxs, + std::vector parquet_schemas, bool output_object_data, + const ParquetFileOptions& parquet_options) : VFileFormatTransformer(state, output_vexpr_ctxs, output_object_data), _parquet_schemas(std::move(parquet_schemas)), - _parquet_options(parquet_options), - _iceberg_schema_json(iceberg_schema_json) { - _iceberg_schema = nullptr; + _parquet_options(parquet_options) { _outstream = std::shared_ptr(new ParquetOutputStream(file_writer)); } -Status VParquetTransformer::_parse_properties() { +Status VParquetWriter::_parse_properties() { try { arrow::MemoryPool* pool = get_arrow_memory_pool(); @@ -241,48 +228,22 @@ Status VParquetTransformer::_parse_properties() { return Status::OK(); } -Status VParquetTransformer::_parse_schema() { - std::vector> fields; - if (_iceberg_schema != nullptr) { - RETURN_IF_ERROR( - iceberg::ArrowSchemaUtil::convert(_iceberg_schema, _state->timezone(), fields)); - } else { - for (size_t i = 0; i < _output_vexpr_ctxs.size(); i++) { - std::shared_ptr type; - RETURN_IF_ERROR(convert_to_arrow_type(_output_vexpr_ctxs[i]->root()->data_type(), &type, - _state->timezone())); - if (!_parquet_schemas.empty()) { - std::shared_ptr field = - arrow::field(_parquet_schemas[i].schema_column_name, type, - _output_vexpr_ctxs[i]->root()->is_nullable()); - fields.emplace_back(field); - } else { - std::shared_ptr field = arrow::field( - _column_names[i], type, _output_vexpr_ctxs[i]->root()->is_nullable()); - fields.emplace_back(field); - } - } - } - - if (_iceberg_schema_json != nullptr) { - std::shared_ptr schema_metadata = - arrow::KeyValueMetadata::Make({"iceberg.schema"}, {*_iceberg_schema_json}); - _arrow_schema = arrow::schema(std::move(fields), std::move(schema_metadata)); - } else { - _arrow_schema = arrow::schema(std::move(fields)); - } - return Status::OK(); +std::unique_ptr VParquetWriter::_create_arrow_block_convertor( + DataTypes types, std::vector names, const std::string& timezone_name, + const cctz::time_zone& timezone) const { + return std::make_unique(std::move(types), std::move(names), + timezone_name, timezone); } -Status VParquetTransformer::write(const Block& block) { +Status VParquetWriter::write(const Block& block) { if (block.rows() == 0) { return Status::OK(); } // serialize std::shared_ptr result; - RETURN_IF_ERROR(convert_to_arrow_batch(block, _arrow_schema, get_arrow_memory_pool(), &result, - _state->timezone_obj())); + RETURN_IF_ERROR( + _arrow_block_convertor->convert_to_arrow(block, get_arrow_memory_pool(), &result)); if (_write_size == 0) { RETURN_DORIS_STATUS_IF_ERROR(_writer->NewBufferedRowGroup()); } @@ -294,16 +255,34 @@ Status VParquetTransformer::write(const Block& block) { return Status::OK(); } -arrow::Status VParquetTransformer::_open_file_writer() { - ARROW_ASSIGN_OR_RAISE(_writer, ::parquet::arrow::FileWriter::Open( - *_arrow_schema, get_arrow_memory_pool(), _outstream, - _parquet_writer_properties, _arrow_properties)); +arrow::Status VParquetWriter::_open_file_writer() { + ARROW_ASSIGN_OR_RAISE(_writer, + ::parquet::arrow::FileWriter::Open( + *_arrow_block_convertor->arrow_schema(), get_arrow_memory_pool(), + _outstream, _parquet_writer_properties, _arrow_properties)); return arrow::Status::OK(); } -Status VParquetTransformer::open() { +Status VParquetWriter::open() { + _timezone = _state->timezone(); + _timezone_obj = _state->timezone_obj(); RETURN_IF_ERROR(_parse_properties()); - RETURN_IF_ERROR(_parse_schema()); + DataTypes types; + types.reserve(_output_vexpr_ctxs.size()); + for (const auto& context : _output_vexpr_ctxs) { + types.emplace_back(context->root()->data_type()); + } + std::vector names = _column_names; + if (!_parquet_schemas.empty()) { + names.clear(); + names.reserve(_parquet_schemas.size()); + for (const auto& schema : _parquet_schemas) { + names.emplace_back(schema.schema_column_name); + } + } + _arrow_block_convertor = _create_arrow_block_convertor(std::move(types), std::move(names), + _timezone, _timezone_obj); + RETURN_IF_ERROR(_arrow_block_convertor->init()); try { RETURN_DORIS_STATUS_IF_ERROR(_open_file_writer()); } catch (const ::parquet::ParquetStatusException& e) { @@ -316,11 +295,11 @@ Status VParquetTransformer::open() { return Status::OK(); } -int64_t VParquetTransformer::written_len() { +int64_t VParquetWriter::written_len() { return _outstream->get_written_len(); } -Status VParquetTransformer::close() { +Status VParquetWriter::close() { try { if (_writer != nullptr) { RETURN_DORIS_STATUS_IF_ERROR(_writer->Close()); @@ -335,85 +314,4 @@ Status VParquetTransformer::close() { return Status::OK(); } -Status VParquetTransformer::collect_file_statistics_after_close(TIcebergColumnStats* stats) { - std::shared_ptr<::parquet::FileMetaData> file_metadata = _writer->metadata(); - if (file_metadata == nullptr) { - return Status::InternalError("File metadata is not available"); - } - std::map column_sizes; - std::map value_counts; - std::map null_value_counts; - std::map lower_bounds; - std::map upper_bounds; - std::map> merged_column_stats; - std::unordered_set variant_field_ids; - - const int num_row_groups = file_metadata->num_row_groups(); - const int num_columns = file_metadata->num_columns(); - for (int col_idx = 0; col_idx < num_columns; ++col_idx) { - const auto& schema_node = file_metadata->schema()->Column(col_idx)->schema_node(); - const auto* parent = schema_node->parent(); - const bool is_variant_child = parent != nullptr && parent->logical_type() != nullptr && - parent->logical_type()->is_variant(); - if (is_variant_child && schema_node->name() != "metadata") { - // The value leaf has no independent Iceberg field and its byte statistics are not - // logical Variant statistics. - continue; - } - const int field_id = is_variant_child ? parent->field_id() : schema_node->field_id(); - if (field_id < 0) { - // Parquet structural leaves (including Variant children) may intentionally omit an - // Iceberg field id. Never publish them under the synthetic -1 key. - continue; - } - if (is_variant_child) { - variant_field_ids.insert(field_id); - } - - for (int rg_idx = 0; rg_idx < num_row_groups; ++rg_idx) { - auto row_group = file_metadata->RowGroup(rg_idx); - auto column_chunk = row_group->ColumnChunk(col_idx); - if (!is_variant_child) { - column_sizes[field_id] += column_chunk->total_compressed_size(); - } - - if (column_chunk->is_stats_set()) { - auto column_stat = column_chunk->statistics(); - if (!merged_column_stats.contains(field_id)) { - merged_column_stats[field_id] = column_stat; - } else { - parquet_utils::merge_stats(merged_column_stats[field_id], column_stat); - } - } - } - } - - bool has_any_null_count = false; - bool has_any_min_max = false; - for (const auto& [field_id, column_stat] : merged_column_stats) { - value_counts[field_id] = column_stat->num_values(); - if (column_stat->HasNullCount()) { - has_any_null_count = true; - int64_t null_count = column_stat->null_count(); - null_value_counts[field_id] = null_count; - value_counts[field_id] += null_count; - } - if (!variant_field_ids.contains(field_id) && column_stat->HasMinMax()) { - has_any_min_max = true; - lower_bounds[field_id] = column_stat->EncodeMin(); - upper_bounds[field_id] = column_stat->EncodeMax(); - } - } - - stats->__set_column_sizes(column_sizes); - stats->__set_value_counts(value_counts); - if (has_any_null_count) { - stats->__set_null_value_counts(null_value_counts); - } - if (has_any_min_max) { - stats->__set_lower_bounds(lower_bounds); - stats->__set_upper_bounds(upper_bounds); - } - return Status::OK(); -} } // namespace doris diff --git a/be/src/format/transformer/vparquet_transformer.h b/be/src/format/transformer/vparquet_writer.h similarity index 72% rename from be/src/format/transformer/vparquet_transformer.h rename to be/src/format/transformer/vparquet_writer.h index e89b9a51435b70..7d7d31a558d04d 100644 --- a/be/src/format/transformer/vparquet_transformer.h +++ b/be/src/format/transformer/vparquet_writer.h @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -28,7 +29,7 @@ #include -#include "format/table/iceberg/schema.h" +#include "format/arrow/arrow_block_convertor.h" #include "format/transformer/vfile_format_transformer.h" namespace doris { @@ -85,23 +86,20 @@ struct ParquetFileOptions { bool enable_int96_timestamps = false; }; -// a wrapper of parquet output stream -class VParquetTransformer final : public VFileFormatTransformer { +// Writes Doris blocks as Parquet files, including schema and Arrow conversion. +class VParquetWriter : public VFileFormatTransformer { public: - VParquetTransformer(RuntimeState* state, doris::io::FileWriter* file_writer, - const VExprContextSPtrs& output_vexpr_ctxs, - std::vector column_names, bool output_object_data, - const ParquetFileOptions& parquet_options, - const std::string* iceberg_schema_json = nullptr, - const iceberg::Schema* iceberg_schema = nullptr); + VParquetWriter(RuntimeState* state, doris::io::FileWriter* file_writer, + const VExprContextSPtrs& output_vexpr_ctxs, + std::vector column_names, bool output_object_data, + const ParquetFileOptions& parquet_options); - VParquetTransformer(RuntimeState* state, doris::io::FileWriter* file_writer, - const VExprContextSPtrs& output_vexpr_ctxs, - std::vector parquet_schemas, bool output_object_data, - const ParquetFileOptions& parquet_options, - const std::string* iceberg_schema_json = nullptr); + VParquetWriter(RuntimeState* state, doris::io::FileWriter* file_writer, + const VExprContextSPtrs& output_vexpr_ctxs, + std::vector parquet_schemas, bool output_object_data, + const ParquetFileOptions& parquet_options); - ~VParquetTransformer() override = default; + ~VParquetWriter() override = default; Status open() override; @@ -111,25 +109,29 @@ class VParquetTransformer final : public VFileFormatTransformer { int64_t written_len() override; - Status collect_file_statistics_after_close(TIcebergColumnStats* stats); +protected: + // Construct the schema and column bindings together for each writer instance. + virtual std::unique_ptr _create_arrow_block_convertor( + DataTypes types, std::vector names, const std::string& timezone_name, + const cctz::time_zone& timezone) const; + std::shared_ptr<::parquet::FileMetaData> _file_metadata() const { return _writer->metadata(); } private: + std::unique_ptr _arrow_block_convertor; Status _parse_properties(); - Status _parse_schema(); arrow::Status _open_file_writer(); std::shared_ptr _outstream; std::shared_ptr<::parquet::WriterProperties> _parquet_writer_properties; std::shared_ptr<::parquet::ArrowWriterProperties> _arrow_properties; std::unique_ptr<::parquet::arrow::FileWriter> _writer; - std::shared_ptr _arrow_schema; std::vector _column_names; std::vector _parquet_schemas; const ParquetFileOptions _parquet_options; - const std::string* _iceberg_schema_json; + std::string _timezone; + cctz::time_zone _timezone_obj; uint64_t _write_size = 0; - const iceberg::Schema* _iceberg_schema; }; } // namespace doris diff --git a/be/src/service/arrow_flight/arrow_flight_batch_reader.cpp b/be/src/service/arrow_flight/arrow_flight_batch_reader.cpp index 8dbc426b882c53..7376128179919a 100644 --- a/be/src/service/arrow_flight/arrow_flight_batch_reader.cpp +++ b/be/src/service/arrow_flight/arrow_flight_batch_reader.cpp @@ -110,8 +110,8 @@ arrow::Status ArrowFlightBatchLocalReader::ReadNextImpl(std::shared_ptr result; cctz::time_zone _timezone_obj; //default UTC - Status stt = convert_to_arrow_batch(*block, block_arrow_schema, - arrow::default_memory_pool(), &result, _timezone_obj); + Status stt = DorisArrowBlockConvertor(block_arrow_schema, _timezone_obj) + .convert_to_arrow(*block, arrow::default_memory_pool(), &result); EXPECT_EQ(Status::OK(), stt) << "convert block to arrow failed" << stt.to_string(); std::cout << "arrow serialize result: " << result->num_columns() << ", " << result->num_rows() << std::endl; diff --git a/be/test/core/data_type_serde/data_type_serde_arrow_test.cpp b/be/test/core/data_type_serde/data_type_serde_arrow_test.cpp index 926850fc38c044..45a66f432db804 100644 --- a/be/test/core/data_type_serde/data_type_serde_arrow_test.cpp +++ b/be/test/core/data_type_serde/data_type_serde_arrow_test.cpp @@ -41,6 +41,7 @@ #include #include +#include #include #include #include @@ -60,6 +61,7 @@ #include "core/column/column_nullable.h" #include "core/column/column_string.h" #include "core/column/column_struct.h" +#include "core/column/column_varbinary.h" #include "core/column/column_vector.h" #include "core/column/variant_v2/column_variant_v2.h" #include "core/data_type/common_data_type_serder_test.h" @@ -79,7 +81,9 @@ #include "core/data_type/data_type_quantilestate.h" #include "core/data_type/data_type_string.h" #include "core/data_type/data_type_struct.h" +#include "core/data_type/data_type_time.h" #include "core/data_type/data_type_timestamptz.h" +#include "core/data_type/data_type_varbinary.h" #include "core/data_type/data_type_variant_v2.h" #include "core/data_type/define_primitive_type.h" #include "core/field.h" @@ -91,6 +95,8 @@ #include "exprs/function/parse/variant_string_parse.h" #include "format/arrow/arrow_block_convertor.h" #include "format/arrow/arrow_row_batch.h" +#include "format/table/iceberg/iceberg_arrow_block_convertor.h" +#include "format/table/paimon/paimon_arrow_block_convertor.h" #include "runtime/descriptors.cpp" #include "util/string_parser.hpp" @@ -137,6 +143,22 @@ std::shared_ptr create_test_block(std::vector cols, int ro ColumnWithTypeAndName type_and_name(vec->get_ptr(), data_type, col_name); block->insert(std::move(type_and_name)); } break; + case TYPE_TINYINT: { + auto vec = ColumnInt8::create(); + for (int i = 0; i < row_num; ++i) { + vec->get_data().push_back(static_cast(i - 3)); + } + block->insert(ColumnWithTypeAndName(vec->get_ptr(), std::make_shared(), + col_name)); + } break; + case TYPE_SMALLINT: { + auto vec = ColumnInt16::create(); + for (int i = 0; i < row_num; ++i) { + vec->get_data().push_back(static_cast(i * 17 - 50)); + } + block->insert(ColumnWithTypeAndName(vec->get_ptr(), std::make_shared(), + col_name)); + } break; case TYPE_INT: if (is_nullable) { { @@ -167,6 +189,38 @@ std::shared_ptr create_test_block(std::vector cols, int ro block->insert(std::move(type_and_name)); } break; + case TYPE_BIGINT: { + auto vec = ColumnInt64::create(); + for (int i = 0; i < row_num; ++i) { + vec->get_data().push_back(static_cast(i) * 1'000'000'007 - 2); + } + block->insert(ColumnWithTypeAndName(vec->get_ptr(), std::make_shared(), + col_name)); + } break; + case TYPE_FLOAT: { + auto vec = ColumnFloat32::create(); + for (int i = 0; i < row_num; ++i) { + vec->get_data().push_back(static_cast(i) + 0.25F); + } + block->insert(ColumnWithTypeAndName(vec->get_ptr(), std::make_shared(), + col_name)); + } break; + case TYPE_DOUBLE: { + auto vec = ColumnFloat64::create(); + for (int i = 0; i < row_num; ++i) { + vec->get_data().push_back(static_cast(i) + 0.125); + } + block->insert(ColumnWithTypeAndName(vec->get_ptr(), std::make_shared(), + col_name)); + } break; + case TYPE_TIMEV2: { + auto vec = ColumnTimeV2::create(); + for (int i = 0; i < row_num; ++i) { + vec->get_data().push_back(3600.125 + i); + } + block->insert(ColumnWithTypeAndName(vec->get_ptr(), std::make_shared(6), + col_name)); + } break; case TYPE_DECIMAL32: { DataTypePtr decimal_data_type = std::make_shared(9, 2); type_desc = decimal_data_type; @@ -230,16 +284,32 @@ std::shared_ptr create_test_block(std::vector cols, int ro col_name); block->insert(type_and_name); } break; - case TYPE_STRING: { + case TYPE_STRING: + case TYPE_VARCHAR: + case TYPE_CHAR: { auto strcol = ColumnString::create(); for (int i = 0; i < row_num; ++i) { std::string is = std::to_string(i); strcol->insert_data(is.c_str(), is.size()); } - DataTypePtr data_type(std::make_shared()); + DataTypePtr data_type = cols[i] == TYPE_STRING + ? std::make_shared() + : std::make_shared( + cols[i] == TYPE_CHAR ? 16 : 128, cols[i]); ColumnWithTypeAndName type_and_name(strcol->get_ptr(), data_type, col_name); block->insert(type_and_name); } break; + case TYPE_VARBINARY: { + auto binary = ColumnVarbinary::create(); + for (int i = 0; i < row_num; ++i) { + const std::array value = {static_cast(i), '\0', + static_cast(0x80 + i), + static_cast(0xff)}; + binary->insert_data(value.data(), value.size()); + } + block->insert(ColumnWithTypeAndName( + binary->get_ptr(), std::make_shared(128), col_name)); + } break; case TYPE_HLL: { DataTypePtr hll_data_type(std::make_shared()); auto hll_column = hll_data_type->create_column(); @@ -253,6 +323,25 @@ std::shared_ptr create_test_block(std::vector cols, int ro block->insert(type_and_name); } break; + case TYPE_BITMAP: { + DataTypePtr bitmap_type(std::make_shared()); + auto bitmap_column = ColumnBitmap::create(); + for (int i = 0; i < row_num; ++i) { + bitmap_column->insert_value(i == 0 ? BitmapValue::empty_bitmap() : BitmapValue(i)); + } + block->insert(ColumnWithTypeAndName(bitmap_column->get_ptr(), bitmap_type, col_name)); + } break; + case TYPE_QUANTILE_STATE: { + DataTypePtr quantile_type(std::make_shared()); + auto quantile_column = ColumnQuantileState::create(); + for (int i = 0; i < row_num; ++i) { + QuantileState state; + state.add_value(i + 0.5); + quantile_column->insert_value(state); + } + block->insert( + ColumnWithTypeAndName(quantile_column->get_ptr(), quantile_type, col_name)); + } break; case TYPE_DATEV2: { auto column_vector_date_v2 = ColumnVector::create(); auto& date_v2_data = column_vector_date_v2->get_data(); @@ -435,6 +524,27 @@ void serialize_and_deserialize_arrow_test(std::vector cols, int r CommonDataTypeSerdeTest::compare_two_blocks(block, assert_block); } +template +std::unique_ptr make_arrow_convertor( + const std::shared_ptr& schema, const cctz::time_zone& timezone) { + return std::make_unique(schema, timezone); +} + +using ArrowConvertorFactory = std::unique_ptr (*)( + const std::shared_ptr&, const cctz::time_zone&); + +// These binding tests supply target schemas intentionally, including unsupported bindings. +// Production table writers construct converters from their native schema parameters instead. +Status convert_to_arrow_batch_for_test( + const Block& block, const std::shared_ptr& schema, arrow::MemoryPool* pool, + std::shared_ptr* result, const cctz::time_zone& timezone, + size_t start = 0, size_t end = 0, + ArrowConvertorFactory factory = make_arrow_convertor) { + auto converter = factory(schema, timezone); + RETURN_IF_ERROR(converter->init()); + return converter->convert_to_arrow(block, pool, result, start, end); +} + void block_converter_test(std::vector cols, int row_num, bool is_nullable) { std::shared_ptr source_block = create_test_block(cols, row_num, is_nullable); std::shared_ptr record_batch; @@ -442,23 +552,27 @@ void block_converter_test(std::vector cols, int row_num, bool is_ Status status = Status::OK(); status = get_arrow_schema_from_block(*source_block, &schema, TimezoneUtils::default_time_zone); ASSERT_TRUE(status.ok() && schema); - cctz::time_zone default_timezone; //default UTC - status = convert_to_arrow_batch(*source_block, schema, arrow::default_memory_pool(), - &record_batch, default_timezone); - ASSERT_TRUE(status.ok() && record_batch); + cctz::time_zone default_timezone; + ASSERT_TRUE( + TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, default_timezone)); + status = convert_to_arrow_batch_for_test(*source_block, schema, arrow::default_memory_pool(), + &record_batch, default_timezone); + ASSERT_TRUE(status.ok() && record_batch) << status; auto target_block = std::make_shared(source_block->clone_empty()); DataTypes source_data_types = source_block->get_data_types(); - status = convert_from_arrow_batch(record_batch, source_data_types, &*target_block, - default_timezone); + status = DorisArrowBlockConvertor(record_batch->schema(), default_timezone) + .convert_from_arrow(record_batch, source_data_types, &*target_block); ASSERT_TRUE(status.ok() && target_block); CommonDataTypeSerdeTest::compare_two_blocks(source_block, target_block); } TEST(DataTypeSerDeArrowTest, DataTypeScalaSerDeTest) { std::vector cols = { - TYPE_INT, TYPE_INT, TYPE_STRING, TYPE_DECIMAL128I, TYPE_BOOLEAN, - TYPE_DECIMAL32, TYPE_DECIMAL64, TYPE_IPV4, TYPE_IPV6, TYPE_LARGEINT, - TYPE_DATETIME, TYPE_DATETIMEV2, TYPE_DATE, TYPE_DATEV2, + TYPE_TINYINT, TYPE_SMALLINT, TYPE_INT, TYPE_BIGINT, TYPE_FLOAT, + TYPE_DOUBLE, TYPE_BOOLEAN, TYPE_STRING, TYPE_VARCHAR, TYPE_CHAR, + TYPE_VARBINARY, TYPE_DECIMAL32, TYPE_DECIMAL64, TYPE_DECIMAL128I, TYPE_IPV4, + TYPE_IPV6, TYPE_LARGEINT, TYPE_DATETIME, TYPE_DATETIMEV2, TYPE_DATE, + TYPE_DATEV2, }; serialize_and_deserialize_arrow_test(cols, 7, true); serialize_and_deserialize_arrow_test(cols, 7, false); @@ -470,6 +584,113 @@ TEST(DataTypeSerDeArrowTest, DataTypeCollectionSerDeTest) { serialize_and_deserialize_arrow_test(cols, 7, false); } +TEST(DataTypeSerDeArrowTest, ArrowBlockConvertorReusesBothDirectionsAndValidatesSlices) { + auto source = create_test_block({TYPE_INT, TYPE_STRING, TYPE_VARBINARY}, 4, true); + std::shared_ptr schema; + ASSERT_TRUE(get_arrow_schema_from_block(*source, &schema, "UTC").ok()); + ArrowFlightArrowBlockConvertor convertor(schema, cctz::utc_time_zone()); + std::shared_ptr batch; + ASSERT_TRUE(convertor.convert_to_arrow(*source, arrow::default_memory_pool(), &batch).ok()); + auto output = std::make_shared(source->clone_empty()); + ASSERT_TRUE(convertor.convert_from_arrow(batch, source->get_data_types(), output.get()).ok()); + CommonDataTypeSerdeTest::compare_two_blocks(source, output); + ASSERT_TRUE( + convertor.convert_to_arrow(*source, arrow::default_memory_pool(), &batch, 1, 3).ok()); + EXPECT_EQ(batch->num_rows(), 2); + EXPECT_FALSE( + convertor.convert_to_arrow(*source, arrow::default_memory_pool(), &batch, 3, 2).ok()); + EXPECT_FALSE( + convertor.convert_to_arrow(*source, arrow::default_memory_pool(), &batch, 5, 0).ok()); +} + +void expect_target_converter_matches_plain(const std::vector& types, + ArrowConvertorFactory target_converter) { + auto block = create_test_block(types, 4, false); + std::shared_ptr schema; + ASSERT_TRUE(get_arrow_schema_from_block(*block, &schema, "UTC").ok()); + + std::shared_ptr plain_batch; + ASSERT_TRUE( + convert_to_arrow_batch_for_test(*block, schema, arrow::default_memory_pool(), + &plain_batch, cctz::utc_time_zone(), 1, block->rows(), + make_arrow_convertor) + .ok()); + std::shared_ptr target_batch; + ASSERT_TRUE(convert_to_arrow_batch_for_test(*block, schema, arrow::default_memory_pool(), + &target_batch, cctz::utc_time_zone(), 1, + block->rows(), target_converter) + .ok()); + ASSERT_TRUE(target_batch->ValidateFull().ok()) << target_batch->ValidateFull(); + EXPECT_TRUE(target_batch->Equals(*plain_batch)); +} + +TEST(DataTypeSerDeArrowTest, IcebergCommonScalarTypesUseDeclaredConverter) { + expect_target_converter_matches_plain( + {TYPE_BOOLEAN, TYPE_INT, TYPE_BIGINT, TYPE_FLOAT, TYPE_DOUBLE, TYPE_STRING, + TYPE_VARBINARY, TYPE_DECIMAL32, TYPE_DECIMAL64, TYPE_DECIMAL128I, TYPE_DATEV2}, + make_arrow_convertor); +} + +TEST(DataTypeSerDeArrowTest, PaimonDoesNotReadThroughGenericArrowSerde) { + Block output; + const auto batch = arrow::RecordBatch::Make(arrow::schema({}), 0, arrow::ArrayVector {}); + paimon::PaimonArrowBlockConvertor converter(batch->schema(), cctz::utc_time_zone()); + const auto status = converter.convert_from_arrow(batch, {}, &output); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(output.columns(), 0); +} + +TEST(DataTypeSerDeArrowTest, PaimonCommonScalarTypesUseDeclaredConverter) { + expect_target_converter_matches_plain( + {TYPE_BOOLEAN, TYPE_TINYINT, TYPE_SMALLINT, TYPE_INT, TYPE_BIGINT, TYPE_FLOAT, + TYPE_DOUBLE, TYPE_STRING, TYPE_VARCHAR, TYPE_CHAR, TYPE_VARBINARY, TYPE_DECIMAL32, + TYPE_DECIMAL64, TYPE_DECIMAL128I, TYPE_DATEV2}, + make_arrow_convertor); +} + +TEST(DataTypeSerDeArrowTest, PlainArrowWritesAggregateStateBinaryTypes) { + auto block = create_test_block({TYPE_HLL, TYPE_BITMAP, TYPE_QUANTILE_STATE}, 3, false); + std::shared_ptr schema; + ASSERT_TRUE(get_arrow_schema_from_block(*block, &schema, "UTC").ok()); + std::shared_ptr batch; + Status status = convert_to_arrow_batch_for_test( + *block, schema, arrow::default_memory_pool(), &batch, cctz::utc_time_zone(), 0, + block->rows(), make_arrow_convertor); + ASSERT_TRUE(status.ok()) << status; + ASSERT_TRUE(batch->ValidateFull().ok()) << batch->ValidateFull(); + ASSERT_EQ(3, batch->num_columns()); + for (const auto& column : batch->columns()) { + EXPECT_EQ(arrow::Type::BINARY, column->type_id()); + EXPECT_EQ(3, column->length()); + } +} + +TEST(DataTypeSerDeArrowTest, PlainArrowWritesTimeV2) { + auto block = create_test_block({TYPE_TIMEV2}, 3, false); + std::shared_ptr schema; + ASSERT_TRUE(get_arrow_schema_from_block(*block, &schema, "UTC").ok()); + ASSERT_EQ(arrow::Type::DOUBLE, schema->field(0)->type()->id()); + + std::shared_ptr batch; + Status status = convert_to_arrow_batch_for_test( + *block, schema, arrow::default_memory_pool(), &batch, cctz::utc_time_zone(), 0, + block->rows(), make_arrow_convertor); + ASSERT_TRUE(status.ok()) << status; + ASSERT_TRUE(batch->ValidateFull().ok()) << batch->ValidateFull(); + const auto& values = assert_cast(*batch->column(0)); + ASSERT_EQ(3, values.length()); + EXPECT_DOUBLE_EQ(3600.125, values.Value(0)); + EXPECT_DOUBLE_EQ(3602.125, values.Value(2)); +} + +TEST(DataTypeSerDeArrowTest, TargetConvertersRecurseThroughOrdinaryComplexTypes) { + const std::vector complex_types = {TYPE_ARRAY, TYPE_MAP, TYPE_STRUCT}; + expect_target_converter_matches_plain( + complex_types, make_arrow_convertor); + expect_target_converter_matches_plain(complex_types, + make_arrow_convertor); +} + TEST(DataTypeSerDeArrowTest, DataTypeMapNullKeySerDeTest) { std::string col_name = "map_null_key"; auto block = std::make_shared(); @@ -535,6 +756,120 @@ TEST(DataTypeSerDeArrowTest, BigStringSerDeTest) { CommonDataTypeSerdeTest::compare_two_blocks(block, assert_block); } +TEST(DataTypeSerDeArrowTest, PaimonTimestampBindsTargetTimezone) { + auto block = create_test_block({TYPE_DATETIMEV2}, 2, false); + auto ntz_schema = + arrow::schema({arrow::field("0", arrow::timestamp(arrow::TimeUnit::MILLI), false)}); + auto ltz_schema = arrow::schema( + {arrow::field("0", arrow::timestamp(arrow::TimeUnit::MILLI, "Asia/Shanghai"), false)}); + cctz::time_zone shanghai; + ASSERT_TRUE(cctz::load_time_zone("Asia/Shanghai", &shanghai)); + + const auto convert = [&](const std::shared_ptr& schema, + ArrowConvertorFactory converter, + std::shared_ptr* record_batch) { + return convert_to_arrow_batch_for_test(*block, schema, arrow::default_memory_pool(), + record_batch, shanghai, 0, block->rows(), converter); + }; + + std::shared_ptr ntz_batch; + Status status = convert(ntz_schema, make_arrow_convertor, + &ntz_batch); + ASSERT_TRUE(status.ok()) << status; + + std::shared_ptr ltz_batch; + status = convert(ltz_schema, make_arrow_convertor, + <z_batch); + ASSERT_TRUE(status.ok()) << status; + const auto& ntz_values = assert_cast(*ntz_batch->column(0)); + const auto& ltz_values = assert_cast(*ltz_batch->column(0)); + EXPECT_EQ(ntz_values.Value(0) - 8 * 60 * 60 * 1000, ltz_values.Value(0)); + + std::shared_ptr iceberg_ntz_batch; + status = convert(ntz_schema, make_arrow_convertor, + &iceberg_ntz_batch); + ASSERT_TRUE(status.ok()) << status; + std::shared_ptr iceberg_ltz_batch; + status = convert(ltz_schema, make_arrow_convertor, + &iceberg_ltz_batch); + ASSERT_TRUE(status.ok()) << status; + EXPECT_TRUE(iceberg_ntz_batch->Equals(*ntz_batch)); + EXPECT_TRUE(iceberg_ltz_batch->Equals(*ltz_batch)); + + std::shared_ptr unused_batch; + status = convert(ntz_schema, make_arrow_convertor, + &unused_batch); + ASSERT_TRUE(status.ok()) << status; + EXPECT_TRUE(unused_batch->Equals(*ntz_batch)); +} + +TEST(DataTypeSerDeArrowTest, TargetConvertersWriteNullableTimestampTz) { + auto values = ColumnTimeStampTz::create(); + TimestampTzValue first; + first.unchecked_set_time(1969, 12, 31, 23, 59, 59, 123456); + TimestampTzValue second; + second.unchecked_set_time(2024, 1, 2, 3, 4, 5, 654321); + values->insert_value(first); + values->insert_value(second); + auto null_map = ColumnUInt8::create(); + null_map->get_data().assign({0, 1}); + + Block block; + block.insert(ColumnWithTypeAndName( + ColumnNullable::create(std::move(values), std::move(null_map)), + make_nullable(std::make_shared(6)), "event_time")); + auto schema = arrow::schema({arrow::field( + "event_time", arrow::timestamp(arrow::TimeUnit::MICRO, "Asia/Shanghai"), true)}); + cctz::time_zone shanghai; + ASSERT_TRUE(cctz::load_time_zone("Asia/Shanghai", &shanghai)); + + const auto convert = [&](ArrowConvertorFactory converter, + std::shared_ptr* batch) { + return convert_to_arrow_batch_for_test(block, schema, arrow::default_memory_pool(), batch, + shanghai, 0, block.rows(), converter); + }; + std::shared_ptr plain_batch; + ASSERT_TRUE(convert(make_arrow_convertor, &plain_batch).ok()); + std::shared_ptr iceberg_batch; + ASSERT_TRUE(convert(make_arrow_convertor, &iceberg_batch) + .ok()); + std::shared_ptr paimon_batch; + ASSERT_TRUE( + convert(make_arrow_convertor, &paimon_batch).ok()); + + EXPECT_TRUE(iceberg_batch->Equals(*plain_batch)); + EXPECT_TRUE(paimon_batch->Equals(*plain_batch)); + const auto& timestamps = assert_cast(*paimon_batch->column(0)); + EXPECT_EQ(-876544, timestamps.Value(0)); + EXPECT_TRUE(timestamps.IsNull(1)); +} + +TEST(DataTypeSerDeArrowTest, PaimonTimestampTzPreservesBothSidesOfDstFold) { + auto values = ColumnTimeStampTz::create(); + for (int hour : {8, 9}) { + TimestampTzValue value; + value.unchecked_set_time(2023, 11, 5, hour, 30, 0, 123456); + values->insert_value(value); + } + Block block; + block.insert(ColumnWithTypeAndName(std::move(values), std::make_shared(6), + "event_time")); + auto schema = arrow::schema({arrow::field( + "event_time", arrow::timestamp(arrow::TimeUnit::MICRO, "America/Los_Angeles"), false)}); + cctz::time_zone timezone; + ASSERT_TRUE(cctz::load_time_zone("America/Los_Angeles", &timezone)); + std::shared_ptr batch; + ASSERT_TRUE(convert_to_arrow_batch_for_test( + block, schema, arrow::default_memory_pool(), &batch, timezone, 0, + block.rows(), make_arrow_convertor) + .ok()); + ASSERT_TRUE(batch->ValidateFull().ok()); + const auto& timestamps = assert_cast(*batch->column(0)); + // Both instants display as 01:30 locally, but must remain one hour apart on the wire. + EXPECT_EQ(1699173000123456LL, timestamps.Value(0)); + EXPECT_EQ(1699176600123456LL, timestamps.Value(1)); +} + TEST(DataTypeSerDeArrowTest, IcebergUuidStringToFixedSizeBinary) { auto block = std::make_shared(); auto strcol = ColumnString::create(); @@ -549,8 +884,9 @@ TEST(DataTypeSerDeArrowTest, IcebergUuidStringToFixedSizeBinary) { std::shared_ptr record_batch; cctz::time_zone default_timezone; - Status status = convert_to_arrow_batch(*block, schema, arrow::default_memory_pool(), - &record_batch, default_timezone); + Status status = convert_to_arrow_batch_for_test( + *block, schema, arrow::default_memory_pool(), &record_batch, default_timezone, 0, + block->rows(), make_arrow_convertor); ASSERT_TRUE(status.ok()) << status; ASSERT_NE(nullptr, record_batch); ASSERT_EQ(2, record_batch->num_rows()); @@ -567,6 +903,23 @@ TEST(DataTypeSerDeArrowTest, IcebergUuidStringToFixedSizeBinary) { EXPECT_EQ(0, std::memcmp(uuid_array->GetValue(1), expected1, sizeof(expected1))); } +TEST(DataTypeSerDeArrowTest, PlainArrowConverterDoesNotInferIcebergUuid) { + Block block; + auto column = ColumnString::create(); + column->insert_data("550e8400-e29b-41d4-a716-446655440000", 36); + block.insert(ColumnWithTypeAndName(column->get_ptr(), std::make_shared(), + "uuid_col")); + auto metadata = arrow::KeyValueMetadata::Make({"originalType"}, {"uuid"}); + auto schema = + arrow::schema({arrow::field("uuid_col", arrow::fixed_size_binary(16), true, metadata)}); + + std::shared_ptr record_batch; + const Status status = convert_to_arrow_batch_for_test( + block, schema, arrow::default_memory_pool(), &record_batch, cctz::utc_time_zone()); + EXPECT_EQ(ErrorCode::INVALID_ARGUMENT, status.code()); + EXPECT_NE(std::string::npos, status.to_string().find("Plain Arrow writer is not bound")); +} + TEST(DataTypeSerDeArrowTest, IcebergVariantExtensionAndParquetSchema) { auto make_variant_column = []() { JsonStringToVariantEncoder encoder({.max_json_key_length = 1024, @@ -614,8 +967,9 @@ TEST(DataTypeSerDeArrowTest, IcebergVariantExtensionAndParquetSchema) { std::shared_ptr record_batch; cctz::time_zone default_timezone; - Status status = convert_to_arrow_batch(block, schema, arrow::default_memory_pool(), - &record_batch, default_timezone); + Status status = convert_to_arrow_batch_for_test( + block, schema, arrow::default_memory_pool(), &record_batch, default_timezone, 0, + block.rows(), make_arrow_convertor); ASSERT_TRUE(status.ok()) << status; ASSERT_NE(nullptr, record_batch); ASSERT_TRUE(record_batch->ValidateFull().ok()); @@ -752,8 +1106,9 @@ TEST(DataTypeSerDeArrowTest, NestedIcebergVariantExtensionsAndParquetSchema) { std::shared_ptr record_batch; cctz::time_zone default_timezone; - Status status = convert_to_arrow_batch(block, schema, arrow::default_memory_pool(), - &record_batch, default_timezone); + Status status = convert_to_arrow_batch_for_test( + block, schema, arrow::default_memory_pool(), &record_batch, default_timezone, 0, + block.rows(), make_arrow_convertor); ASSERT_TRUE(status.ok()) << status; ASSERT_NE(nullptr, record_batch); ASSERT_TRUE(record_batch->ValidateFull().ok()) << record_batch->ValidateFull(); @@ -832,8 +1187,9 @@ TEST(DataTypeSerDeArrowTest, NestedIcebergUuidStringToFixedSizeBinary) { std::shared_ptr record_batch; cctz::time_zone default_timezone; - Status status = convert_to_arrow_batch(*block, schema, arrow::default_memory_pool(), - &record_batch, default_timezone); + Status status = convert_to_arrow_batch_for_test( + *block, schema, arrow::default_memory_pool(), &record_batch, default_timezone, 0, + block->rows(), make_arrow_convertor); ASSERT_TRUE(status.ok()) << status; auto struct_array = std::static_pointer_cast(record_batch->column(0)); @@ -843,6 +1199,167 @@ TEST(DataTypeSerDeArrowTest, NestedIcebergUuidStringToFixedSizeBinary) { EXPECT_EQ(0, std::memcmp(uuid_array->GetValue(0), expected, sizeof(expected))); } +TEST(DataTypeSerDeArrowTest, IcebergFixedVarbinaryPreservesRawBytesNullsAndRowRange) { + constexpr int width = 256; + std::vector values(4, std::string(width, '\0')); + for (size_t row = 0; row < values.size(); ++row) { + for (int byte = 0; byte < width; ++byte) { + values[row][byte] = static_cast((row * 67 + byte * 131) & 0xff); + } + } + + auto data = ColumnVarbinary::create(); + for (const auto& value : values) { + data->insert_data(value.data(), value.size()); + } + auto null_map = ColumnUInt8::create(); + null_map->get_data().assign({0, 0, 1, 0}); + auto column = ColumnNullable::create(std::move(data), std::move(null_map)); + DataTypePtr type = make_nullable(std::make_shared(width)); + + Block block; + block.insert(ColumnWithTypeAndName(column->get_ptr(), type, "fixed_col")); + auto schema = arrow::schema({arrow::field("fixed_col", arrow::fixed_size_binary(width), true)}); + + std::shared_ptr record_batch; + Status status = convert_to_arrow_batch_for_test( + block, schema, arrow::default_memory_pool(), &record_batch, cctz::utc_time_zone(), 1, 4, + make_arrow_convertor); + ASSERT_TRUE(status.ok()) << status; + ASSERT_EQ(3, record_batch->num_rows()); + auto fixed = std::static_pointer_cast(record_batch->column(0)); + ASSERT_EQ(width, fixed->byte_width()); + EXPECT_FALSE(fixed->IsNull(0)); + EXPECT_TRUE(fixed->IsNull(1)); + EXPECT_FALSE(fixed->IsNull(2)); + EXPECT_EQ(0, std::memcmp(fixed->GetValue(0), values[1].data(), width)); + EXPECT_EQ(0, std::memcmp(fixed->GetValue(2), values[3].data(), width)); +} + +TEST(DataTypeSerDeArrowTest, IcebergFixedBinaryPreservesBindingsAndRejectsInvalidValues) { + auto convert = [](DataTypePtr type, std::string_view value, int target_width, + ArrowConvertorFactory converter) { + MutableColumnPtr column = type->create_column(); + column->insert_data(value.data(), value.size()); + Block block; + block.insert(ColumnWithTypeAndName(std::move(column), type, "fixed_col")); + auto schema = arrow::schema( + {arrow::field("fixed_col", arrow::fixed_size_binary(target_width), true)}); + std::shared_ptr record_batch; + return convert_to_arrow_batch_for_test(block, schema, arrow::default_memory_pool(), + &record_batch, cctz::utc_time_zone(), 0, + block.rows(), converter); + }; + + const auto& iceberg_converter = make_arrow_convertor; + Status status = convert(std::make_shared(4), "abc", 4, iceberg_converter); + EXPECT_EQ(ErrorCode::INVALID_ARGUMENT, status.code()); + EXPECT_NE(std::string::npos, + status.to_string().find("Fixed size binary column expects 4 bytes, got 3")); + + status = convert(std::make_shared(4), "abcde", 4, iceberg_converter); + EXPECT_EQ(ErrorCode::INVALID_ARGUMENT, status.code()); + EXPECT_NE(std::string::npos, + status.to_string().find("Fixed size binary column expects 4 bytes, got 5")); + + // This refactor preserves the current external mappings and validates physical values. + status = convert(std::make_shared(8), "abcd", 4, iceberg_converter); + EXPECT_TRUE(status.ok()) << status; + status = convert(std::make_shared(4, TYPE_CHAR), "abcd", 4, iceberg_converter); + EXPECT_TRUE(status.ok()) << status; + + status = convert(std::make_shared(4), "abcd", 4, + make_arrow_convertor); + EXPECT_EQ(ErrorCode::INVALID_ARGUMENT, status.code()); + EXPECT_NE(std::string::npos, status.to_string().find("Plain Arrow writer is not bound")); +} + +TEST(DataTypeSerDeArrowTest, NestedIcebergFixedVarbinaryUsesIcebergConverterRecursively) { + constexpr int width = 4; + const std::array values = {std::string("\0\x01\xfe\xff", width), + std::string("abcd", width), + std::string("\x80\0\x7f\x10", width)}; + DataTypePtr fixed_type = std::make_shared(width); + DataTypePtr nullable_fixed_type = make_nullable(fixed_type); + + auto make_nullable_fixed_column = [&]() { + auto data = ColumnVarbinary::create(); + for (const auto& value : values) { + data->insert_data(value.data(), value.size()); + } + auto null_map = ColumnUInt8::create(); + null_map->get_data().assign({0, 1, 0}); + return ColumnNullable::create(std::move(data), std::move(null_map)); + }; + + auto array_offsets = ColumnArray::ColumnOffsets::create(); + array_offsets->get_data().assign({2, 3, 3}); + auto array_column = ColumnArray::create(make_nullable_fixed_column(), std::move(array_offsets)); + DataTypePtr array_type = std::make_shared(nullable_fixed_type); + + auto map_keys_data = ColumnString::create(); + for (std::string_view key : {"k0", "k1", "k2"}) { + map_keys_data->insert_data(key.data(), key.size()); + } + auto map_key_nulls = ColumnUInt8::create(); + map_key_nulls->get_data().assign({0, 0, 0}); + auto map_keys = ColumnNullable::create(std::move(map_keys_data), std::move(map_key_nulls)); + auto map_offsets = ColumnArray::ColumnOffsets::create(); + map_offsets->get_data().assign({2, 3, 3}); + auto map_column = ColumnMap::create(std::move(map_keys), make_nullable_fixed_column(), + std::move(map_offsets)); + DataTypePtr map_type = std::make_shared( + make_nullable(std::make_shared()), nullable_fixed_type); + + MutableColumns struct_children; + struct_children.emplace_back(make_nullable_fixed_column()); + auto struct_column = ColumnStruct::create(std::move(struct_children)); + DataTypePtr struct_type = + std::make_shared(DataTypes {nullable_fixed_type}, Strings {"payload"}); + + Block block; + block.insert(ColumnWithTypeAndName(std::move(array_column), array_type, "items")); + block.insert(ColumnWithTypeAndName(std::move(map_column), map_type, "attrs")); + block.insert(ColumnWithTypeAndName(std::move(struct_column), struct_type, "info")); + + const auto arrow_fixed = arrow::fixed_size_binary(width); + auto schema = arrow::schema({ + arrow::field("items", arrow::list(arrow::field("element", arrow_fixed, true)), true), + arrow::field("attrs", + std::make_shared(arrow::field("key", arrow::utf8(), false), + arrow::field("value", arrow_fixed, true)), + true), + arrow::field("info", arrow::struct_({arrow::field("payload", arrow_fixed, true)}), + true), + }); + + std::shared_ptr record_batch; + Status status = convert_to_arrow_batch_for_test( + block, schema, arrow::default_memory_pool(), &record_batch, cctz::utc_time_zone(), 0, + block.rows(), make_arrow_convertor); + ASSERT_TRUE(status.ok()) << status; + ASSERT_TRUE(record_batch->ValidateFull().ok()) << record_batch->ValidateFull(); + + auto items = std::static_pointer_cast(record_batch->column(0)); + auto item_values = std::static_pointer_cast(items->values()); + EXPECT_EQ(3, item_values->length()); + EXPECT_TRUE(item_values->IsNull(1)); + EXPECT_EQ(0, std::memcmp(item_values->GetValue(0), values[0].data(), width)); + EXPECT_EQ(0, std::memcmp(item_values->GetValue(2), values[2].data(), width)); + + auto attrs = std::static_pointer_cast(record_batch->column(1)); + auto attr_values = std::static_pointer_cast(attrs->items()); + EXPECT_EQ(3, attr_values->length()); + EXPECT_TRUE(attr_values->IsNull(1)); + EXPECT_EQ(0, std::memcmp(attr_values->GetValue(2), values[2].data(), width)); + + auto info = std::static_pointer_cast(record_batch->column(2)); + auto payloads = std::static_pointer_cast(info->field(0)); + EXPECT_EQ(3, payloads->length()); + EXPECT_TRUE(payloads->IsNull(1)); + EXPECT_EQ(0, std::memcmp(payloads->GetValue(0), values[0].data(), width)); +} + TEST(DataTypeSerDeArrowTest, CharToFixedSizeBinaryPadsZeros) { auto block = std::make_shared(); auto strcol = ColumnString::create(); @@ -854,8 +1371,9 @@ TEST(DataTypeSerDeArrowTest, CharToFixedSizeBinaryPadsZeros) { std::shared_ptr record_batch; cctz::time_zone default_timezone; - Status status = convert_to_arrow_batch(*block, schema, arrow::default_memory_pool(), - &record_batch, default_timezone); + Status status = convert_to_arrow_batch_for_test( + *block, schema, arrow::default_memory_pool(), &record_batch, default_timezone, 0, + block->rows(), make_arrow_convertor); ASSERT_TRUE(status.ok()) << status; auto fixed_array = @@ -875,8 +1393,8 @@ TEST(DataTypeSerDeArrowTest, StringToLargeBinary) { std::shared_ptr record_batch; cctz::time_zone default_timezone; - Status status = convert_to_arrow_batch(*block, schema, arrow::default_memory_pool(), - &record_batch, default_timezone); + Status status = convert_to_arrow_batch_for_test(*block, schema, arrow::default_memory_pool(), + &record_batch, default_timezone); ASSERT_TRUE(status.ok()) << status; auto binary_array = std::static_pointer_cast(record_batch->column(0)); @@ -926,6 +1444,16 @@ TEST(DataTypeSerDeArrowTest, ConvertDateTimeV2ToNaiveArrowType) { EXPECT_TRUE(timestamp_type->timezone().empty()); } +TEST(DataTypeSerDeArrowTest, CanonicalizeUtcTimezoneForArrow) { + const auto timestamptz_type = std::make_shared(6); + std::shared_ptr arrow_type; + + auto status = convert_to_arrow_type(timestamptz_type, &arrow_type, "Z", true); + ASSERT_TRUE(status.ok()) << status; + const auto timestamp_type = std::static_pointer_cast(arrow_type); + EXPECT_EQ("UTC", timestamp_type->timezone()); +} + TEST(DataTypeSerDeArrowTest, DateTimeV2ArrowEncodingFollowsSchemaTimezone) { auto datetime_column = ColumnVector::create(); DateV2Value datetime_value; @@ -940,8 +1468,8 @@ TEST(DataTypeSerDeArrowTest, DateTimeV2ArrowEncodingFollowsSchemaTimezone) { auto timezone_schema = arrow::schema( {arrow::field("ts", arrow::timestamp(arrow::TimeUnit::MICRO, "+08:00"), false)}); std::shared_ptr timezone_batch; - auto status = convert_to_arrow_batch(block, timezone_schema, arrow::default_memory_pool(), - &timezone_batch, utc_plus_eight); + auto status = convert_to_arrow_batch_for_test( + block, timezone_schema, arrow::default_memory_pool(), &timezone_batch, utc_plus_eight); ASSERT_TRUE(status.ok()) << status; const auto timezone_type = std::static_pointer_cast( @@ -955,8 +1483,8 @@ TEST(DataTypeSerDeArrowTest, DateTimeV2ArrowEncodingFollowsSchemaTimezone) { auto naive_schema = arrow::schema({arrow::field("ts", arrow::timestamp(arrow::TimeUnit::MICRO), false)}); std::shared_ptr naive_batch; - status = convert_to_arrow_batch(block, naive_schema, arrow::default_memory_pool(), &naive_batch, - utc_plus_eight); + status = convert_to_arrow_batch_for_test(block, naive_schema, arrow::default_memory_pool(), + &naive_batch, utc_plus_eight); ASSERT_TRUE(status.ok()) << status; const auto naive_type = @@ -968,4 +1496,50 @@ TEST(DataTypeSerDeArrowTest, DateTimeV2ArrowEncodingFollowsSchemaTimezone) { EXPECT_EQ(1783004400123456, naive_array->Value(0)); } +TEST(DataTypeSerDeArrowTest, NestedDateTimeV2PlainArrowAcceptsNaiveSchema) { + DateV2Value datetime_value; + datetime_value.unchecked_set_time(2026, 7, 2, 15, 0, 0, 123456); + const auto datetime_type = std::make_shared(6); + + const auto array_type = std::make_shared(datetime_type); + auto array_column = array_type->create_column(); + Array array_value; + array_value.push_back(Field::create_field(datetime_value)); + array_column->insert(Field::create_field(array_value)); + + // ColumnMap stores nullable key/value children even though Arrow rejects null map keys. + const auto map_type = std::make_shared( + make_nullable(std::make_shared()), make_nullable(datetime_type)); + auto map_column = map_type->create_column(); + Array map_keys; + map_keys.push_back(Field::create_field("event")); + Array map_values; + map_values.push_back(Field::create_field(datetime_value)); + Map map_value; + map_value.push_back(Field::create_field(map_keys)); + map_value.push_back(Field::create_field(map_values)); + map_column->insert(Field::create_field(map_value)); + + const auto struct_type = + std::make_shared(DataTypes {datetime_type}, Strings {"event_time"}); + auto struct_column = struct_type->create_column(); + Struct struct_value; + struct_value.push_back(Field::create_field(datetime_value)); + struct_column->insert(Field::create_field(struct_value)); + + Block block; + block.insert(ColumnWithTypeAndName(array_column->get_ptr(), array_type, "events")); + block.insert(ColumnWithTypeAndName(map_column->get_ptr(), map_type, "event_map")); + block.insert(ColumnWithTypeAndName(struct_column->get_ptr(), struct_type, "event_struct")); + + std::shared_ptr naive_schema; + ASSERT_TRUE(get_arrow_schema_from_block(block, &naive_schema, "Asia/Shanghai", true).ok()); + std::shared_ptr naive_batch; + const auto status = convert_to_arrow_batch_for_test( + block, naive_schema, arrow::default_memory_pool(), &naive_batch, + cctz::fixed_time_zone(std::chrono::hours(8))); + ASSERT_TRUE(status.ok()) << status; + ASSERT_TRUE(naive_batch->ValidateFull().ok()) << naive_batch->ValidateFull(); +} + } // namespace doris diff --git a/be/test/core/data_type_serde/data_type_variant_v2_serde_output_test.cpp b/be/test/core/data_type_serde/data_type_variant_v2_serde_output_test.cpp index c895bb5a87506b..25c3ca7abd9aa9 100644 --- a/be/test/core/data_type_serde/data_type_variant_v2_serde_output_test.cpp +++ b/be/test/core/data_type_serde/data_type_variant_v2_serde_output_test.cpp @@ -32,19 +32,23 @@ #include "core/arena.h" #include "core/assert_cast.h" +#include "core/column/column_array.h" #include "core/column/column_const.h" #include "core/column/column_nullable.h" #include "core/column/column_string.h" #include "core/column/column_vector.h" #include "core/column/variant_v2/column_variant_v2.h" +#include "core/data_type/data_type_array.h" #include "core/data_type/data_type_date_or_datetime_v2.h" #include "core/data_type/data_type_number.h" #include "core/data_type/data_type_string.h" +#include "core/data_type/data_type_variant_v2.h" #include "core/data_type_serde/data_type_nullable_serde.h" #include "core/data_type_serde/data_type_variant_v2_serde.h" #include "core/string_buffer.hpp" #include "core/value/variant/variant_parquet_encoding.h" #include "exprs/function/parse/variant_string_parse.h" +#include "format/table/paimon/paimon_arrow_block_convertor.h" #include "util/mysql_row_buffer.h" namespace doris { @@ -213,12 +217,16 @@ std::unique_ptr binary_variant_arrow_builder() { std::make_shared(arrow::default_memory_pool())}); } -void expect_binary_variant_bytes(const DataTypeVariantV2SerDe& serde, const IColumn& column, - const ColumnVariantV2& encoded, +void expect_paimon_variant_bytes(const IColumn& column, const ColumnVariantV2& encoded, const NullMap* null_map = nullptr) { auto builder = binary_variant_arrow_builder(); - const Status status = serde.write_column_to_arrow(column, null_map, builder.get(), 0, - column.size(), cctz::utc_time_zone()); + DataTypePtr type = std::make_shared(); + const auto serde = type->get_serde(); + const auto field = arrow::field("payload", binary_variant_arrow_type(), true); + const Status status = + paimon::PaimonArrowBlockConvertor(arrow::schema({field}), cctz::utc_time_zone()) + .write_column(type, *serde, column, null_map, field, builder.get(), 0, + column.size(), cctz::utc_time_zone()); ASSERT_TRUE(status.ok()) << status; std::shared_ptr output; @@ -443,21 +451,64 @@ TEST(DataTypeVariantV2SerdeOutputTest, ConstNullableAndOuterMasksPreserveBoundar EXPECT_TRUE(invalid_dates->is_typed()); } -TEST(DataTypeVariantV2SerdeOutputTest, BinaryStructPreservesEncodedAndTypedBytesAndOuterNulls) { - DataTypeVariantV2SerDe serde; +TEST(PaimonArrowBlockConvertorTest, BinaryStructPreservesEncodedAndTypedBytesAndOuterNulls) { auto documents = encoded_json({R"({"a":[1,null,"x"]})", R"({"hidden":true})", "null"}); NullMap mask {0, 1, 0}; - expect_binary_variant_bytes(serde, *documents, *documents, &mask); + expect_paimon_variant_bytes(*documents, *documents, &mask); auto typed = typed_strings( {std::string_view("plain"), std::nullopt, std::string_view(R"({"text":"value"})")}); ColumnPtr encoded = encoded_copy(*typed); - expect_binary_variant_bytes(serde, *typed, assert_cast(*encoded)); + expect_paimon_variant_bytes(*typed, assert_cast(*encoded)); EXPECT_TRUE(typed->is_typed()); } -TEST(DataTypeVariantV2SerdeOutputTest, BinaryStructRejectsUnsupportedPaimonPrimitive) { - DataTypeVariantV2SerDe serde; +TEST(PaimonArrowBlockConvertorTest, NestedArrayUsesPaimonSerdeRecursively) { + auto variants = encoded_json({R"({"id":1})", R"({"hidden":true})", R"([true,"x"])", "null"}); + const auto expected = variants->read_view(); + auto offsets = ColumnArray::ColumnOffsets::create(); + offsets->get_data().assign({2, 4}); + // DataTypeArray always wraps its element type in Nullable, including VARIANT. + auto nulls = ColumnUInt8::create(4, 0); + nulls->get_data()[1] = 1; + auto elements = ColumnNullable::create(std::move(variants), std::move(nulls)); + auto array = ColumnArray::create(std::move(elements), std::move(offsets)); + + DataTypePtr variant_type = std::make_shared(); + DataTypePtr array_type = std::make_shared(variant_type); + auto arrow_type = arrow::list(binary_variant_arrow_type()); + auto field = arrow::field("payloads", arrow_type, false); + std::unique_ptr builder; + ASSERT_TRUE(arrow::MakeBuilder(arrow::default_memory_pool(), arrow_type, &builder).ok()); + + const auto serde = array_type->get_serde(); + Status status = paimon::PaimonArrowBlockConvertor(arrow::schema({field}), cctz::utc_time_zone()) + .write_column(array_type, *serde, *array, nullptr, field, builder.get(), + 0, array->size(), cctz::utc_time_zone()); + ASSERT_TRUE(status.ok()) << status; + + std::shared_ptr output; + ASSERT_TRUE(builder->Finish(&output).ok()); + const auto& lists = assert_cast(*output); + const auto& structs = assert_cast(*lists.values()); + const auto& values = assert_cast(*structs.field(0)); + const auto& metadata = assert_cast(*structs.field(1)); + ASSERT_TRUE(output->ValidateFull().ok()); + ASSERT_EQ(4, structs.length()); + EXPECT_TRUE(structs.IsNull(1)); + EXPECT_FALSE(structs.IsNull(3)); + for (size_t row : {0, 2, 3}) { + const VariantRef value = expected.value_at(row); + const auto actual_value = values.GetView(row); + const auto actual_metadata = metadata.GetView(row); + EXPECT_EQ(std::string_view(actual_value.data(), actual_value.size()), + std::string_view(value.value.data, value.value.size)); + EXPECT_EQ(std::string_view(actual_metadata.data(), actual_metadata.size()), + std::string_view(value.metadata.data, value.metadata.size)); + } +} + +TEST(PaimonArrowBlockConvertorTest, BinaryStructRejectsUnsupportedPaimonPrimitive) { VariantBatchBuilder builder(VariantBatchBuilder::ReserveHint {.rows = 1}); auto row = builder.begin_row(); row.add_time_ntz_micros(1'500'000); @@ -465,8 +516,13 @@ TEST(DataTypeVariantV2SerdeOutputTest, BinaryStructRejectsUnsupportedPaimonPrimi auto encoded = ColumnVariantV2::create(); encoded->insert_encoded_batch(builder.finish_batch()); auto arrow_builder = binary_variant_arrow_builder(); - const Status status = serde.write_column_to_arrow(*encoded, nullptr, arrow_builder.get(), 0, - encoded->size(), cctz::utc_time_zone()); + DataTypePtr type = std::make_shared(); + const auto serde = type->get_serde(); + const auto field = arrow::field("payload", binary_variant_arrow_type(), true); + const Status status = + paimon::PaimonArrowBlockConvertor(arrow::schema({field}), cctz::utc_time_zone()) + .write_column(type, *serde, *encoded, nullptr, field, arrow_builder.get(), 0, + encoded->size(), cctz::utc_time_zone()); EXPECT_EQ(status.code(), ErrorCode::NOT_IMPLEMENTED_ERROR); EXPECT_NE(status.to_string().find("Paimon does not support Variant primitive id 17"), std::string::npos); diff --git a/be/test/format/arrow/arrow_block_convertor_test.cpp b/be/test/format/arrow/arrow_block_convertor_test.cpp new file mode 100644 index 00000000000000..635b790b0f866c --- /dev/null +++ b/be/test/format/arrow/arrow_block_convertor_test.cpp @@ -0,0 +1,291 @@ +// 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/arrow/arrow_block_convertor.h" + +#include +#include +#include +#include +#include + +#include "core/column/column_vector.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_factory.hpp" +#include "core/data_type/data_type_map.h" +#include "core/data_type/data_type_nullable.h" +#include "core/data_type/data_type_number.h" +#include "core/data_type/data_type_struct.h" +#include "format/parquet/parquet_arrow_block_convertor.h" +#include "format/table/hive/hive_arrow_block_convertor.h" +#include "format/table/iceberg/iceberg_arrow_block_convertor.h" +#include "format/table/iceberg/schema.h" +#include "format/table/iceberg/schema_parser.h" +#include "format/table/paimon/paimon_arrow_block_convertor.h" +#include "udf/python/python_udf_meta.h" +#include "util/timezone_utils.h" + +namespace doris { + +class ArrowBlockConvertorTest : public testing::Test { +protected: + static void SetUpTestSuite() { TimezoneUtils::load_timezones_to_cache(); } +}; + +TEST_F(ArrowBlockConvertorTest, ParquetOwnsSchemaAndTimezoneParameters) { + cctz::time_zone shanghai; + ASSERT_TRUE(TimezoneUtils::find_cctz_time_zone("Asia/Shanghai", shanghai)); + DataTypes types {DataTypeFactory::instance().create_data_type(TYPE_DATETIMEV2, false, 0, 6), + DataTypeFactory::instance().create_data_type(TYPE_TIMESTAMPTZ, false, 0, 6)}; + ParquetArrowBlockConvertor parquet(types, {"local_time", "instant"}, "Asia/Shanghai", shanghai); + hive::HiveArrowBlockConvertor hive(types, {"local_time", "instant"}, "UTC", + cctz::utc_time_zone()); + ASSERT_TRUE(parquet.init().ok()); + ASSERT_TRUE(hive.init().ok()); + const auto timestamp = [](const ArrowBlockConvertor& converter, + int index) -> const arrow::TimestampType& { + return static_cast( + *converter.arrow_schema()->field(index)->type()); + }; + EXPECT_EQ("Asia/Shanghai", timestamp(parquet, 0).timezone()); + EXPECT_EQ("Asia/Shanghai", timestamp(parquet, 1).timezone()); + EXPECT_EQ("UTC", timestamp(hive, 0).timezone()); + // A second writer must not change an existing writer's schema or timestamp binding. + ASSERT_TRUE(hive.init().ok()); + EXPECT_EQ("Asia/Shanghai", timestamp(parquet, 0).timezone()); + EXPECT_EQ("Asia/Shanghai", timestamp(parquet, 1).timezone()); +} + +TEST_F(ArrowBlockConvertorTest, ParquetRejectsMismatchedColumnNames) { + ParquetArrowBlockConvertor converter({std::make_shared()}, {}, "UTC", + cctz::utc_time_zone()); + EXPECT_FALSE(converter.init().ok()); + EXPECT_EQ(nullptr, converter.arrow_schema()); +} + +TEST_F(ArrowBlockConvertorTest, IcebergBuildsItsOwnSchemaAndMetadata) { + const std::string json = + R"({"type":"struct","fields":[{"id":7,"name":"payload","required":false,"type":"variant"}]})"; + auto schema = iceberg::SchemaParser::from_json(json); + iceberg::IcebergArrowBlockConvertor converter(*schema, &json, "UTC", cctz::utc_time_zone()); + ASSERT_TRUE(converter.init().ok()); + const auto& arrow_schema = converter.arrow_schema(); + ASSERT_NE(nullptr, arrow_schema); + ASSERT_EQ(arrow::Type::EXTENSION, arrow_schema->field(0)->type()->id()); + const auto& variant = static_cast(*arrow_schema->field(0)->type()); + EXPECT_EQ(arrow::Type::STRUCT, variant.storage_type()->id()); + EXPECT_EQ("arrow.parquet.variant", variant.extension_name()); + ASSERT_NE(nullptr, arrow_schema->metadata()); + EXPECT_EQ(json, arrow_schema->metadata()->Get("iceberg.schema").ValueOrDie()); + EXPECT_EQ("7", arrow_schema->field(0)->metadata()->Get("PARQUET:field_id").ValueOrDie()); +} + +TEST_F(ArrowBlockConvertorTest, PaimonDecodesPinnedSchemaAndValidatesBlock) { + auto schema = arrow::schema({arrow::field( + "items", + arrow::list(arrow::field("item", arrow::timestamp(arrow::TimeUnit::MICRO), false)), + false)}); + auto sink = arrow::io::BufferOutputStream::Create().ValueOrDie(); + auto writer = arrow::ipc::MakeStreamWriter(sink, schema).ValueOrDie(); + ASSERT_TRUE(writer->Close().ok()); + auto buffer = sink->Finish().ValueOrDie(); + paimon::PaimonArrowBlockConvertor converter(buffer->ToString(), cctz::utc_time_zone()); + ASSERT_TRUE(converter.init().ok()); + EXPECT_TRUE(schema->Equals(*converter.arrow_schema(), true)); + Block block; + std::shared_ptr batch; + EXPECT_FALSE(converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch).ok()); + EXPECT_EQ(nullptr, batch); +} + +TEST_F(ArrowBlockConvertorTest, PaimonRejectsInvalidSerializedSchema) { + for (const std::string& bytes : {std::string(), std::string("invalid schema")}) { + paimon::PaimonArrowBlockConvertor converter(bytes, cctz::utc_time_zone()); + EXPECT_FALSE(converter.init().ok()); + EXPECT_EQ(nullptr, converter.arrow_schema()); + } +} + +TEST_F(ArrowBlockConvertorTest, PythonBuildsSchemaAndConvertsSlicesInternally) { + auto type = std::make_shared(); + auto column = ColumnInt32::create(); + column->get_data().assign({11, 22, 33}); + Block block; + block.insert({std::move(column), type, "value"}); + PythonArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone()); + std::shared_ptr batch; + EXPECT_FALSE(converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch).ok()); + ASSERT_TRUE(converter.init().ok()); + EXPECT_EQ("value", converter.arrow_schema()->field(0)->name()); + ASSERT_TRUE(converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch, 1, 3).ok()); + ASSERT_TRUE(batch->ValidateFull().ok()); + ASSERT_EQ(2, batch->num_rows()); + const auto& values = static_cast(*batch->column(0)); + EXPECT_EQ(22, values.Value(0)); + EXPECT_EQ(33, values.Value(1)); + Block output; + ASSERT_TRUE(converter.convert_from_arrow(batch, {type}, &output).ok()); + EXPECT_EQ(2, output.rows()); +} + +TEST_F(ArrowBlockConvertorTest, PaimonSerializedSchemaKeepsTimestampBindingsPerInstance) { + cctz::time_zone shanghai; + ASSERT_TRUE(cctz::load_time_zone("Asia/Shanghai", &shanghai)); + auto type = DataTypeFactory::instance().create_data_type(TYPE_DATETIMEV2, false, 0, 6); + DateV2Value datetime; + const std::string format = "%Y-%m-%d %H:%i:%s.%f"; + const std::string value = "1969-12-31 23:59:59.999999"; + ASSERT_TRUE(datetime.from_date_format_str(format.data(), format.size(), value.data(), + value.size())); + auto column = ColumnDateTimeV2::create(); + column->insert_value(datetime); + Block block; + block.insert({std::move(column), type, "ts"}); + std::vector> converters; + for (const std::string& zone : {std::string(), std::string("Asia/Shanghai")}) { + auto schema = arrow::schema( + {arrow::field("ts", arrow::timestamp(arrow::TimeUnit::MICRO, zone), false)}); + auto sink = arrow::io::BufferOutputStream::Create().ValueOrDie(); + auto writer = arrow::ipc::MakeStreamWriter(sink, schema).ValueOrDie(); + ASSERT_TRUE(writer->Close().ok()); + auto converter = std::make_unique( + sink->Finish().ValueOrDie()->ToString(), shanghai); + ASSERT_TRUE(converter->init().ok()); + converters.emplace_back(std::move(converter)); + } + // Both converters remain alive while writing, so accidental shared schema state is observable. + for (size_t i = 0; i < converters.size(); ++i) { + std::shared_ptr batch; + ASSERT_TRUE( + converters[i]->convert_to_arrow(block, arrow::default_memory_pool(), &batch).ok()); + ASSERT_TRUE(batch->ValidateFull().ok()); + const auto& values = static_cast(*batch->column(0)); + EXPECT_EQ(i == 0 ? -1LL : -28800000001LL, values.Value(0)); + } +} + +TEST_F(ArrowBlockConvertorTest, PythonFixedOffsetSchemaMatchesDeclaredProtocol) { + auto type = DataTypeFactory::instance().create_data_type(TYPE_DATETIMEV2, false, 0, 6); + DataTypes types {type, std::make_shared(type)}; + Block block; + for (size_t i = 0; i < types.size(); ++i) { + block.insert({types[i]->create_column(), types[i], "arg" + std::to_string(i)}); + } + for (const std::string& zone : + {TimezoneUtils::default_time_zone, std::string("+05:45"), std::string("-03:30"), + std::string("UTC"), std::string("Asia/Shanghai")}) { + SCOPED_TRACE(zone); + cctz::time_zone timezone; + ASSERT_TRUE(TimezoneUtils::find_cctz_time_zone(zone, timezone)); + PythonArrowBlockConvertor converter(block, zone, timezone); + ASSERT_TRUE(converter.init().ok()); + std::shared_ptr declared; + ASSERT_TRUE(PythonUDFMeta::convert_types_to_schema(types, zone, &declared).ok()); + EXPECT_TRUE(declared->Equals(*converter.arrow_schema())) + << "declared=" << declared->ToString() + << ", actual=" << converter.arrow_schema()->ToString(); + } +} + +TEST_F(ArrowBlockConvertorTest, PythonFixedOffsetBatchesPreserveValuesAndNulls) { + auto type = make_nullable( + DataTypeFactory::instance().create_data_type(TYPE_DATETIMEV2, false, 0, 6)); + DateV2Value value; + value.unchecked_set_time(2023, 4, 20, 0, 0, 0, 123456); + auto column = type->create_column(); + column->insert(Field::create_field(value)); + column->insert_default(); + Block block; + block.insert({std::move(column), type, "arg0"}); + for (const auto& [zone, offset_seconds] : std::vector> { + {"+08:00", 28800}, {"+05:45", 20700}, {"-03:30", -12600}, {"UTC", 0}}) { + SCOPED_TRACE(zone); + cctz::time_zone timezone; + ASSERT_TRUE(TimezoneUtils::find_cctz_time_zone(zone, timezone)); + PythonArrowBlockConvertor converter(block, zone, timezone); + ASSERT_TRUE(converter.init().ok()); + std::shared_ptr batch; + ASSERT_TRUE(converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch).ok()); + ASSERT_TRUE(batch->ValidateFull().ok()); + std::shared_ptr declared; + ASSERT_TRUE(PythonUDFMeta::convert_types_to_schema({type}, zone, &declared).ok()); + EXPECT_TRUE(declared->Equals(*batch->schema())); + const auto& values = static_cast(*batch->column(0)); + // The protocol label and the epoch must describe the same wall-clock input. + EXPECT_EQ(1681948800123456LL - offset_seconds * 1000000, values.Value(0)); + EXPECT_TRUE(values.IsNull(1)); + Block output; + ASSERT_TRUE(converter.convert_from_arrow(batch, {type}, &output).ok()); + const auto& actual = *output.get_by_position(0).column; + const auto& expected = *block.get_by_position(0).column; + ASSERT_EQ(expected.size(), actual.size()); + EXPECT_EQ(0, expected.compare_at(0, 0, actual, 1)); + EXPECT_TRUE(actual.is_null_at(1)); + } +} + +TEST_F(ArrowBlockConvertorTest, TableWritersPreserveFixedOffsetSchemaNames) { + auto type = DataTypeFactory::instance().create_data_type(TYPE_TIMESTAMPTZ, false, 0, 6); + const std::string json = + R"({"type":"struct","fields":[{"id":1,"name":"ts","required":true,"type":"timestamptz"}]})"; + auto schema = iceberg::SchemaParser::from_json(json); + for (const std::string& zone : {std::string("+05:45"), std::string("-03:30")}) { + SCOPED_TRACE(zone); + cctz::time_zone timezone; + ASSERT_TRUE(TimezoneUtils::find_cctz_time_zone(zone, timezone)); + ParquetArrowBlockConvertor parquet({type}, {"ts"}, zone, timezone); + hive::HiveArrowBlockConvertor hive({type}, {"ts"}, zone, timezone); + iceberg::IcebergArrowBlockConvertor iceberg(*schema, &json, zone, timezone); + for (ArrowBlockConvertor* converter : + {static_cast(&parquet), static_cast(&hive), + static_cast(&iceberg)}) { + ASSERT_TRUE(converter->init().ok()); + const auto& timestamp = static_cast( + *converter->arrow_schema()->field(0)->type()); + EXPECT_EQ(zone, timestamp.timezone()); + } + } +} + +TEST_F(ArrowBlockConvertorTest, TableConvertersRejectMismatchedNestedSchemas) { + auto integer = std::make_shared(); + DataTypes types {std::make_shared(integer), + std::make_shared(make_nullable(integer), make_nullable(integer)), + std::make_shared(DataTypes {integer}, Strings {"value"})}; + for (const auto& type : types) { + SCOPED_TRACE(type->get_name()); + Block block; + auto column = type->create_column(); + column->insert_default(); + block.insert({std::move(column), type, "nested"}); + for (const auto& target : {arrow::int32(), arrow::struct_({})}) { + auto schema = arrow::schema({arrow::field("nested", target)}); + paimon::PaimonArrowBlockConvertor paimon(schema, cctz::utc_time_zone()); + iceberg::IcebergArrowBlockConvertor iceberg(schema, cctz::utc_time_zone()); + for (ArrowBlockConvertor* converter : {static_cast(&paimon), + static_cast(&iceberg)}) { + std::shared_ptr batch; + const auto status = + converter->convert_to_arrow(block, arrow::default_memory_pool(), &batch); + EXPECT_EQ(ErrorCode::INVALID_ARGUMENT, status.code()) << status; + EXPECT_EQ(nullptr, batch); + } + } + } +} + +} // namespace doris diff --git a/be/test/format/transformer/vparquet_transformer_test.cpp b/be/test/format/transformer/vparquet_writer_test.cpp similarity index 61% rename from be/test/format/transformer/vparquet_transformer_test.cpp rename to be/test/format/transformer/vparquet_writer_test.cpp index 9a93adb7c55d3d..00446c9fb43bc2 100644 --- a/be/test/format/transformer/vparquet_transformer_test.cpp +++ b/be/test/format/transformer/vparquet_writer_test.cpp @@ -15,10 +15,13 @@ // specific language governing permissions and limitations // under the License. -#include "format/transformer/vparquet_transformer.h" +#include "format/transformer/vparquet_writer.h" +#include +#include #include #include +#include #include #include @@ -28,19 +31,24 @@ #include "core/column/column_nullable.h" #include "core/column/variant_v2/column_variant_v2.h" #include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_factory.hpp" #include "core/data_type/data_type_nullable.h" #include "core/data_type/data_type_variant_v2.h" #include "exprs/function/parse/variant_string_parse.h" #include "format/table/iceberg/schema_parser.h" +#include "format/transformer/viceberg_parquet_writer.h" #include "io/fs/local_file_system.h" #include "runtime/runtime_state.h" #include "testutil/mock/mock_slot_ref.h" +#include "util/timezone_utils.h" #include "util/uid_util.h" namespace doris { -class VParquetTransformerTest : public testing::Test { +class VParquetWriterTest : public testing::Test { protected: + static void SetUpTestSuite() { TimezoneUtils::load_timezones_to_cache(); } + void SetUp() override { _file_path = "./vparquet_transformer_" + UniqueId::gen_uid().to_string() + ".parquet"; _fs = io::global_local_filesystem(); @@ -52,7 +60,7 @@ class VParquetTransformerTest : public testing::Test { std::shared_ptr _fs; }; -TEST_F(VParquetTransformerTest, WritesIcebergVariantAndCollectsLogicalMetrics) { +TEST_F(VParquetWriterTest, WritesIcebergVariantAndCollectsLogicalMetrics) { auto variant_type = std::make_shared(); auto nullable_variant_type = make_nullable(variant_type); VExprContextSPtrs output_exprs = @@ -74,8 +82,8 @@ TEST_F(VParquetTransformerTest, WritesIcebergVariantAndCollectsLogicalMetrics) { .parquet_version = TParquetVersion::PARQUET_1_0, .parquet_disable_dictionary = false, .enable_int96_timestamps = false}; - VParquetTransformer transformer(&state, file_writer.get(), output_exprs, {"payload"}, false, - options, &schema_json, schema.get()); + VIcebergParquetWriter transformer(&state, file_writer.get(), output_exprs, {"payload"}, false, + options, &schema_json, *schema); ASSERT_TRUE(transformer.open().ok()); JsonStringToVariantEncoder encoder({.max_json_key_length = 1024, @@ -128,7 +136,7 @@ TEST_F(VParquetTransformerTest, WritesIcebergVariantAndCollectsLogicalMetrics) { EXPECT_EQ(-1, payload_group.field(1)->field_id()); } -TEST_F(VParquetTransformerTest, WritesNestedIcebergVariant) { +TEST_F(VParquetWriterTest, WritesNestedIcebergVariant) { auto variant_type = std::make_shared(); auto array_type = std::make_shared(variant_type); auto nullable_array_type = make_nullable(array_type); @@ -161,8 +169,8 @@ TEST_F(VParquetTransformerTest, WritesNestedIcebergVariant) { .parquet_version = TParquetVersion::PARQUET_1_0, .parquet_disable_dictionary = false, .enable_int96_timestamps = false}; - VParquetTransformer transformer(&state, file_writer.get(), output_exprs, {"events"}, false, - options, &schema_json, schema.get()); + VIcebergParquetWriter transformer(&state, file_writer.get(), output_exprs, {"events"}, false, + options, &schema_json, *schema); ASSERT_TRUE(transformer.open().ok()); JsonStringToVariantEncoder encoder({.max_json_key_length = 1024, @@ -210,4 +218,88 @@ TEST_F(VParquetTransformerTest, WritesNestedIcebergVariant) { EXPECT_EQ(3, element->field_id()); } +TEST_F(VParquetWriterTest, PreservesInt64TimestampSchema) { + DataTypes types {DataTypeFactory::instance().create_data_type(TYPE_DATETIMEV2, false, 0, 6), + DataTypeFactory::instance().create_data_type(TYPE_TIMESTAMPTZ, false, 0, 6)}; + VExprContextSPtrs output_exprs = MockSlotRef::create_mock_contexts(types); + + io::FileWriterPtr file_writer; + ASSERT_TRUE(_fs->create_file(_file_path, &file_writer).ok()); + RuntimeState state; + state.set_timezone("Asia/Shanghai"); + ParquetFileOptions options {.compression_type = TParquetCompressionType::UNCOMPRESSED, + .parquet_version = TParquetVersion::PARQUET_1_0, + .parquet_disable_dictionary = false, + .enable_int96_timestamps = false}; + VParquetWriter transformer(&state, file_writer.get(), output_exprs, + std::vector {"local_time", "instant"}, false, options); + ASSERT_TRUE(transformer.open().ok()); + ASSERT_TRUE(transformer.close().ok()); + + auto reader = ::parquet::ParquetFileReader::OpenFile(_file_path, false); + const auto* root = reader->metadata()->schema()->group_node(); + ASSERT_EQ(2, root->field_count()); + ASSERT_NE(nullptr, root->field(0)->logical_type()); + ASSERT_NE(nullptr, root->field(1)->logical_type()); + // Moving schema construction must preserve the current Parquet timestamp representation. + EXPECT_NE(std::string::npos, + root->field(0)->logical_type()->ToString().find("isAdjustedToUTC=true")); + EXPECT_NE(std::string::npos, + root->field(1)->logical_type()->ToString().find("isAdjustedToUTC=true")); +} + +TEST_F(VParquetWriterTest, WritesInt96DatetimeUsingWriterTimezone) { + auto datetime_type = DataTypeFactory::instance().create_data_type(TYPE_DATETIMEV2, false, 0, 6); + VExprContextSPtrs output_exprs = MockSlotRef::create_mock_contexts(DataTypes {datetime_type}); + + io::FileWriterPtr file_writer; + ASSERT_TRUE(_fs->create_file(_file_path, &file_writer).ok()); + RuntimeState state; + state.set_timezone("Asia/Shanghai"); + ParquetFileOptions options {.compression_type = TParquetCompressionType::UNCOMPRESSED, + .parquet_version = TParquetVersion::PARQUET_1_0, + .parquet_disable_dictionary = false, + .enable_int96_timestamps = true}; + VParquetWriter transformer(&state, file_writer.get(), output_exprs, {"local_time"}, false, + options); + ASSERT_TRUE(transformer.open().ok()); + + DateV2Value datetime; + const std::string format = "%Y-%m-%d %H:%i:%s.%f"; + const std::string value = "2023-04-20 00:00:00.123456"; + ASSERT_TRUE(datetime.from_date_format_str(format.data(), format.size(), value.data(), + value.size())); + auto column = ColumnDateTimeV2::create(); + column->insert_value(datetime); + Block block; + block.insert(ColumnWithTypeAndName(std::move(column), datetime_type, "local_time")); + const auto write_status = transformer.write(block); + ASSERT_TRUE(write_status.ok()) << write_status.to_string(); + ASSERT_TRUE(transformer.close().ok()); + + auto physical_reader = ::parquet::ParquetFileReader::OpenFile(_file_path, false); + const auto* root = physical_reader->metadata()->schema()->group_node(); + ASSERT_EQ(1, root->field_count()); + const auto& primitive = assert_cast(*root->field(0)); + EXPECT_EQ(::parquet::Type::INT96, primitive.physical_type()); + + auto input_result = arrow::io::ReadableFile::Open(_file_path); + ASSERT_TRUE(input_result.ok()) << input_result.status(); + auto arrow_reader_result = + ::parquet::arrow::OpenFile(*input_result, arrow::default_memory_pool()); + ASSERT_TRUE(arrow_reader_result.ok()) << arrow_reader_result.status(); + std::unique_ptr<::parquet::arrow::FileReader> arrow_reader = std::move(*arrow_reader_result); + std::shared_ptr table; + ASSERT_TRUE(arrow_reader->ReadTable(&table).ok()); + ASSERT_EQ(1, table->num_rows()); + const auto& timestamp = assert_cast(*table->column(0)->chunk(0)); + const auto& timestamp_type = assert_cast(*timestamp.type()); + int64_t epoch_micros = timestamp.Value(0); + if (timestamp_type.unit() == arrow::TimeUnit::NANO) { + epoch_micros /= 1000; + } + // Hive-compatible INT96 encodes the UTC instant for the writer's local DATETIMEV2 value. + EXPECT_EQ(1681920000123456LL, epoch_micros); +} + } // namespace doris diff --git a/regression-test/data/pythonudf_p0/test_pythonudf_base_data_type.out b/regression-test/data/pythonudf_p0/test_pythonudf_base_data_type.out index 781ea3168073e9..4e3e5e77c11014 100644 --- a/regression-test/data/pythonudf_p0/test_pythonudf_base_data_type.out +++ b/regression-test/data/pythonudf_p0/test_pythonudf_base_data_type.out @@ -1,15 +1,15 @@ -- This file is automatically generated. You should know what you did if you want to edit this -- !select_1 -- -True,127,32767,2147483647,9223372036854775807,170141183460469231731687303715884105727,1.2300000190734863,4.56789,123456.780000000,12345678901.230000000,123456789012345678901.234567890,2023-01-01,2023-01-01 20:34:56+08:00,char_data_1,varchar_data_1,string_data_1 -False,-128,-32768,-2147483648,-9223372036854775808,-170141183460469231731687303715884105728,-2.3399999141693115,-5.6789,-987654.320000000,-98765432.110000000,-987654321098765432.109876540,2024-05-15,2024-05-15 16:22:10+08:00,char_data_2,varchar_data_2,string_data_2 -True,0,0,0,0,0,0.0,0.0,0E-9,0E-9,0E-9,2025-10-15,2025-10-15 08:00:00+08:00,char_zero,varchar_zero,string_zero -False,100,20000,300000000,4000000000000000000,99999999999999999999999999999999999999,3.140000104904175,2.71828,999999.990000000,99999999999999.990000000,100000000000000000000000.000000000,2022-12-31,2023-01-01 07:59:59+08:00,char_max,varchar_max,string_max -True,-50,-10000,-100000000,-5000000000000000000,-99999999999999999999999999999999999999,-1.409999966621399,-0.57721,-0.010000000,-0.010000000,0E-9,2021-07-04,2021-07-04 22:30:00+08:00,char_neg,varchar_neg,string_neg +True,127,32767,2147483647,9223372036854775807,170141183460469231731687303715884105727,1.2300000190734863,4.56789,123456.780000000,12345678901.230000000,123456789012345678901.234567890,2023-01-01,2023-01-01 12:34:56+08:00,char_data_1,varchar_data_1,string_data_1 +False,-128,-32768,-2147483648,-9223372036854775808,-170141183460469231731687303715884105728,-2.3399999141693115,-5.6789,-987654.320000000,-98765432.110000000,-987654321098765432.109876540,2024-05-15,2024-05-15 08:22:10+08:00,char_data_2,varchar_data_2,string_data_2 +True,0,0,0,0,0,0.0,0.0,0E-9,0E-9,0E-9,2025-10-15,2025-10-15 00:00:00+08:00,char_zero,varchar_zero,string_zero +False,100,20000,300000000,4000000000000000000,99999999999999999999999999999999999999,3.140000104904175,2.71828,999999.990000000,99999999999999.990000000,100000000000000000000000.000000000,2022-12-31,2022-12-31 23:59:59+08:00,char_max,varchar_max,string_max +True,-50,-10000,-100000000,-5000000000000000000,-99999999999999999999999999999999999999,-1.409999966621399,-0.57721,-0.010000000,-0.010000000,0E-9,2021-07-04,2021-07-04 14:30:00+08:00,char_neg,varchar_neg,string_neg -- !select_2 -- -True,127,32767,2147483647,9223372036854775807,170141183460469231731687303715884105727,1.2300000190734863,4.56789,123456.780000000,12345678901.230000000,123456789012345678901.234567890,2023-01-01,2023-01-01 20:34:56+08:00,char_data_1,varchar_data_1,string_data_1 -False,-128,-32768,-2147483648,-9223372036854775808,-170141183460469231731687303715884105728,-2.3399999141693115,-5.6789,-987654.320000000,-98765432.110000000,-987654321098765432.109876540,2024-05-15,2024-05-15 16:22:10+08:00,char_data_2,varchar_data_2,string_data_2 -True,0,0,0,0,0,0.0,0.0,0E-9,0E-9,0E-9,2025-10-15,2025-10-15 08:00:00+08:00,char_zero,varchar_zero,string_zero -False,100,20000,300000000,4000000000000000000,99999999999999999999999999999999999999,3.140000104904175,2.71828,999999.990000000,99999999999999.990000000,100000000000000000000000.000000000,2022-12-31,2023-01-01 07:59:59+08:00,char_max,varchar_max,string_max -True,-50,-10000,-100000000,-5000000000000000000,-99999999999999999999999999999999999999,-1.409999966621399,-0.57721,-0.010000000,-0.010000000,0E-9,2021-07-04,2021-07-04 22:30:00+08:00,char_neg,varchar_neg,string_neg +True,127,32767,2147483647,9223372036854775807,170141183460469231731687303715884105727,1.2300000190734863,4.56789,123456.780000000,12345678901.230000000,123456789012345678901.234567890,2023-01-01,2023-01-01 12:34:56+08:00,char_data_1,varchar_data_1,string_data_1 +False,-128,-32768,-2147483648,-9223372036854775808,-170141183460469231731687303715884105728,-2.3399999141693115,-5.6789,-987654.320000000,-98765432.110000000,-987654321098765432.109876540,2024-05-15,2024-05-15 08:22:10+08:00,char_data_2,varchar_data_2,string_data_2 +True,0,0,0,0,0,0.0,0.0,0E-9,0E-9,0E-9,2025-10-15,2025-10-15 00:00:00+08:00,char_zero,varchar_zero,string_zero +False,100,20000,300000000,4000000000000000000,99999999999999999999999999999999999999,3.140000104904175,2.71828,999999.990000000,99999999999999.990000000,100000000000000000000000.000000000,2022-12-31,2022-12-31 23:59:59+08:00,char_max,varchar_max,string_max +True,-50,-10000,-100000000,-5000000000000000000,-99999999999999999999999999999999999999,-1.409999966621399,-0.57721,-0.010000000,-0.010000000,0E-9,2021-07-04,2021-07-04 14:30:00+08:00,char_neg,varchar_neg,string_neg diff --git a/regression-test/data/pythonudf_p0/test_pythonudf_data_types.out b/regression-test/data/pythonudf_p0/test_pythonudf_data_types.out index d6921742e95707..9bf6b418170abd 100644 --- a/regression-test/data/pythonudf_p0/test_pythonudf_data_types.out +++ b/regression-test/data/pythonudf_p0/test_pythonudf_data_types.out @@ -15,7 +15,7 @@ 2024-01-15 -- !select_datetime -- -2024-01-15 18:30:45+08:00 +2024-01-15 10:30:45+08:00 -- !select_table_types -- 1 11 200 1010000 diff --git a/regression-test/suites/pythonudf_p0/test_python_arrow_convertor_timezone.groovy b/regression-test/suites/pythonudf_p0/test_python_arrow_convertor_timezone.groovy new file mode 100644 index 00000000000000..0112e61ebb5a33 --- /dev/null +++ b/regression-test/suites/pythonudf_p0/test_python_arrow_convertor_timezone.groovy @@ -0,0 +1,88 @@ +// 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. + +suite("test_python_arrow_convertor_timezone") { + def runtimeVersion = getPythonUdfRuntimeVersion() + def originalTimezone = sql("select @@time_zone")[0][0] + def scalar = "python_arrow_convertor_scalar" + def tableFunction = "python_arrow_convertor_rows" + def aggregate = "python_arrow_convertor_max" + try { + sql "DROP TABLE IF EXISTS python_arrow_convertor_values" + sql """CREATE TABLE python_arrow_convertor_values (id INT, ts DATETIME(6)) + DISTRIBUTED BY HASH(id) BUCKETS 1 PROPERTIES("replication_num"="1")""" + sql """INSERT INTO python_arrow_convertor_values VALUES + (1, '1969-12-31 23:59:59.999999'), + (2, '2023-04-20 00:00:00.123456'), (3, NULL)""" + [scalar, tableFunction, aggregate].each { name -> + sql "DROP FUNCTION IF EXISTS ${name}(DATETIME(6))" + } + sql """CREATE FUNCTION ${scalar}(DATETIME(6)) RETURNS STRING + PROPERTIES("type"="PYTHON_UDF", "symbol"="evaluate", + "runtime_version"="${runtimeVersion}") AS \$\$ +def evaluate(value): + return None if value is None else value.strftime('%Y-%m-%d %H:%M:%S.%f') +\$\$""" + // Declare the yielded string as the array element; ARRAY> emits a struct. + sql """CREATE TABLES FUNCTION ${tableFunction}(DATETIME(6)) + RETURNS ARRAY + PROPERTIES("type"="PYTHON_UDF", "symbol"="evaluate", + "runtime_version"="${runtimeVersion}") AS \$\$ +def evaluate(value): + yield (None if value is None else value.strftime('%Y-%m-%d %H:%M:%S.%f'),) +\$\$""" + sql """CREATE AGGREGATE FUNCTION ${aggregate}(DATETIME(6)) RETURNS DATETIME(6) + PROPERTIES("type"="PYTHON_UDF", "symbol"="Maximum", + "runtime_version"="${runtimeVersion}") AS \$\$ +class Maximum: + def __init__(self): + self.value = None + @property + def aggregate_state(self): + return self.value + def accumulate(self, value): + if value is not None and (self.value is None or value > self.value): + self.value = value + def merge(self, value): + self.accumulate(value) + def finish(self): + return self.value +\$\$""" + // The Python protocol uses its declared default offset regardless of the session zone. + // Check values observed inside Python as well as returned values to detect offset cancellation. + ["UTC", "+05:45", "-03:30", "Asia/Shanghai"].each { zone -> + sql "SET time_zone = '${zone}'" + assertEquals([[true], [true], [true]], sql(""" + SELECT ${scalar}(ts) <=> date_format(ts, '%Y-%m-%d %H:%i:%s.%f') + FROM python_arrow_convertor_values ORDER BY id""")) + assertEquals([[true], [true], [true]], sql(""" + SELECT result.value <=> date_format(ts, '%Y-%m-%d %H:%i:%s.%f') + FROM python_arrow_convertor_values + LATERAL VIEW ${tableFunction}(ts) result AS value ORDER BY id""")) + assertEquals([[true]], sql(""" + SELECT ${aggregate}(ts) <=> max(ts) FROM python_arrow_convertor_values""")) + assertEquals([[true], [true]], sql(""" + SELECT ${aggregate}(ts) <=> max(ts) FROM python_arrow_convertor_values + GROUP BY id % 2 ORDER BY id % 2""")) + } + } finally { + sql "SET time_zone = '${originalTimezone}'" + [scalar, tableFunction, aggregate].each { name -> + try_sql("DROP FUNCTION IF EXISTS ${name}(DATETIME(6))") + } + } +} diff --git a/regression-test/suites/pythonudf_p0/test_pythonudf_data_types.groovy b/regression-test/suites/pythonudf_p0/test_pythonudf_data_types.groovy index da423263633466..6b15d9382bd953 100644 --- a/regression-test/suites/pythonudf_p0/test_pythonudf_data_types.groovy +++ b/regression-test/suites/pythonudf_p0/test_pythonudf_data_types.groovy @@ -139,6 +139,7 @@ def evaluate(dt): \$\$; """ + // The Arrow timezone label must not shift a DATETIME's civil fields in Python. qt_select_datetime """ SELECT py_datetime_test('2024-01-15 10:30:45') AS result; """ // Test 7: Comprehensive test - create table and test multiple data types