From c070334ead50dd3d8e485c4033134791116e2b84 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 22 Sep 2026 16:46:26 +0800 Subject: [PATCH 1/2] [refactor](arrow) Make format convertors explicit and schema-owned ### What problem does this PR solve? Related PR: #67784. This is the second split, based on the primitives merged in #68297. Arrow batch conversion mixes protocol serialization with table-specific UUID and Variant handling, while writers construct schemas separately. Introduce explicit Doris, Python, Arrow Flight, Parquet, Hive, Iceberg and Paimon convertors, instantiated with their own schema parameters and timezone. Move table schema construction/decoding into the convertors and route nested SerDe writes through the selected format. Separate Parquet, Hive and Iceberg writers and migrate all callers. Preserve current Parquet timestamp encoding, UUID/Variant layouts and external type mappings. Keep Python's numeric conversion timezone consistent with its declared Arrow schema, and reject incompatible nested target schemas before casts or child access. No FE, Thrift, ORC timestamp, or binary computation changes are included. ### Release note Fix Python UDF timestamp conversion to preserve wall-clock values when the Arrow protocol declares a fixed-offset timezone. ### Check List (For Author) - Test - [x] Unit Test: schema ownership/isolation, slices, invalid schemas, nested/null values, UUID/Variant bytes, fixed-offset Python batches, and Parquet timestamp representation. - [x] Regression test: add Python UDF/UDTF/UDAF timezone assertions and extract the corresponding existing Python snapshot corrections from #67784. - Behavior changed: - [x] Yes: align Python UDF conversion with its Arrow timezone declaration; report invalid nested schema bindings as errors. - Does this need documentation? - [x] No. Validation: ASAN BE build and 396 selected tests passed; the final incremental build and 285 focused tests also passed. clang-format 16 and Groovy/embedded Python syntax checks passed. Python SQL and external-catalog end-to-end regressions were not run locally and remain for CI. ### Check List (For Reviewer who merge this PR) - [ ] Confirm the release note - [ ] Confirm test cases - [ ] Confirm document - [ ] Add branch pick label --- .../data_type_serde/data_type_array_serde.cpp | 71 ++ .../data_type_serde/data_type_array_serde.h | 10 + .../data_type_serde/data_type_map_serde.cpp | 103 +++ .../data_type_serde/data_type_map_serde.h | 10 + .../data_type_nullable_serde.cpp | 23 + .../data_type_nullable_serde.h | 10 + be/src/core/data_type_serde/data_type_serde.h | 18 + .../data_type_string_serde.cpp | 47 ++ .../data_type_serde/data_type_string_serde.h | 5 + .../data_type_struct_serde.cpp | 71 ++ .../data_type_serde/data_type_struct_serde.h | 10 + .../data_type_variant_v2_serde.cpp | 396 +++++------ .../data_type_variant_v2_serde.h | 10 + .../operator/memory_scratch_sink_operator.cpp | 9 +- .../iceberg/viceberg_delete_file_writer.cpp | 6 +- .../iceberg/viceberg_partition_writer.cpp | 8 +- .../paimon/jni_paimon_write_backend.cpp | 57 +- .../writer/paimon/jni_paimon_write_backend.h | 11 +- .../exec/sink/writer/vfile_result_writer.cpp | 4 +- .../sink/writer/vhive_partition_writer.cpp | 4 +- .../aggregate_function_python_udaf.cpp | 25 +- be/src/exprs/function/function_python_udf.cpp | 24 +- .../table_function/python_udtf_function.cpp | 14 +- be/src/format/arrow/arrow_block_convertor.cpp | 329 +++++---- be/src/format/arrow/arrow_block_convertor.h | 137 ++-- be/src/format/arrow/arrow_row_batch.cpp | 5 +- .../parquet/parquet_arrow_block_convertor.cpp | 55 ++ .../parquet/parquet_arrow_block_convertor.h | 46 ++ .../table/hive/hive_arrow_block_convertor.h | 31 + .../iceberg/iceberg_arrow_block_convertor.cpp | 57 ++ .../iceberg/iceberg_arrow_block_convertor.h | 51 ++ .../paimon/paimon_arrow_block_convertor.cpp | 58 ++ .../paimon/paimon_arrow_block_convertor.h | 47 ++ .../vfile_format_transformer_factory.cpp | 4 +- .../format/transformer/vhive_parquet_writer.h | 38 ++ .../transformer/viceberg_parquet_writer.cpp | 133 ++++ .../transformer/viceberg_parquet_writer.h | 49 ++ .../transformer/vjni_format_transformer.h | 2 +- be/src/format/transformer/vorc_transformer.h | 2 +- ...et_transformer.cpp => vparquet_writer.cpp} | 196 ++---- ...arquet_transformer.h => vparquet_writer.h} | 42 +- .../arrow_flight_batch_reader.cpp | 8 +- .../data_type/common_data_type_serder_test.h | 4 +- .../data_type_serde_arrow_test.cpp | 628 +++++++++++++++++- ...data_type_variant_v2_serde_output_test.cpp | 80 ++- .../arrow/arrow_block_convertor_test.cpp | 291 ++++++++ ...rmer_test.cpp => vparquet_writer_test.cpp} | 108 ++- .../test_pythonudf_base_data_type.out | 20 +- .../test_pythonudf_data_types.out | 2 +- ...est_python_arrow_convertor_timezone.groovy | 87 +++ .../test_pythonudf_data_types.groovy | 1 + 51 files changed, 2738 insertions(+), 719 deletions(-) create mode 100644 be/src/format/parquet/parquet_arrow_block_convertor.cpp create mode 100644 be/src/format/parquet/parquet_arrow_block_convertor.h create mode 100644 be/src/format/table/hive/hive_arrow_block_convertor.h create mode 100644 be/src/format/table/iceberg/iceberg_arrow_block_convertor.cpp create mode 100644 be/src/format/table/iceberg/iceberg_arrow_block_convertor.h create mode 100644 be/src/format/table/paimon/paimon_arrow_block_convertor.cpp create mode 100644 be/src/format/table/paimon/paimon_arrow_block_convertor.h create mode 100644 be/src/format/transformer/vhive_parquet_writer.h create mode 100644 be/src/format/transformer/viceberg_parquet_writer.cpp create mode 100644 be/src/format/transformer/viceberg_parquet_writer.h rename be/src/format/transformer/{vparquet_transformer.cpp => vparquet_writer.cpp} (54%) rename be/src/format/transformer/{vparquet_transformer.h => vparquet_writer.h} (72%) create mode 100644 be/test/format/arrow/arrow_block_convertor_test.cpp rename be/test/format/transformer/{vparquet_transformer_test.cpp => vparquet_writer_test.cpp} (61%) create mode 100644 regression-test/suites/pythonudf_p0/test_python_arrow_convertor_timezone.groovy 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..5c7db25b7f8afb --- /dev/null +++ b/regression-test/suites/pythonudf_p0/test_python_arrow_convertor_timezone.groovy @@ -0,0 +1,87 @@ +// 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') +\$\$""" + 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 From c552895e6c9dbdece9ff4bfab484a7086cd4555e Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 22 Sep 2026 22:33:37 +0800 Subject: [PATCH 2/2] [fix](regression) Declare the Python timezone UDTF string output ### What problem does this PR solve? Related PR: #68381 The timezone regression declares ARRAY> but compares its lateral-view output directly to a string. Declare ARRAY to match the single string yielded per row and allow the comparison to reach execution. Preserve coverage for fixed-offset timezones, microseconds, pre-epoch values and NULLs. ### Release note None ### Check List (For Author) - Test: Groovy syntax and embedded Python scalar, UDTF and UDAF value checks; SQL analysis verification. - Behavior changed: No (test-only correction). - Does this need documentation: No. --- .../pythonudf_p0/test_python_arrow_convertor_timezone.groovy | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 index 5c7db25b7f8afb..0112e61ebb5a33 100644 --- a/regression-test/suites/pythonudf_p0/test_python_arrow_convertor_timezone.groovy +++ b/regression-test/suites/pythonudf_p0/test_python_arrow_convertor_timezone.groovy @@ -37,8 +37,9 @@ suite("test_python_arrow_convertor_timezone") { 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> + RETURNS ARRAY PROPERTIES("type"="PYTHON_UDF", "symbol"="evaluate", "runtime_version"="${runtimeVersion}") AS \$\$ def evaluate(value):