This is an automated email from the ASF dual-hosted git repository. yiguolei pushed a commit to branch branch-4.2 in repository https://gitbox.apache.org/repos/asf/doris.git
commit 1bbc51ef3c5a891c5b7de9e55d7731e2ba1ec85d Author: Gabriel <[email protected]> AuthorDate: Thu Oct 8 10:46:33 2026 +0800 [fix](arrow-flight) Preserve nested logical type metadata (#68614) ### What problem does this PR solve? Flight SQL attaches `doris_type` metadata to top-level LARGEINT/IP fields, but ARRAY items, MAP keys/values and STRUCT children lose it. A nested LARGEINT therefore arrives as an Arrow string that clients cannot distinguish from ordinary text. JSON and VARIANT also lack the marker at the top level. For example, in `SELECT named_struct('number', CAST(17 AS LARGEINT), 'text', '17')`, both children currently appear as unannotated strings. This change annotates only the LARGEINT child, allowing clients to restore its logical type safely. Related PR: #67530 (closed while awaiting type-mapping refactoring). This change targets the current `branch-4.1` schema and converter implementation. Move schema construction behind the abstract `ArrowSchemaConvertor` interface. Doris, Flight, legacy Flight, Parquet, Iceberg and Paimon own their conversion policies; recursive Doris type conversion no longer accepts protocol or metadata flags. Flight selects its compatibility implementation once during result-sink preparation. Iceberg retains table field IDs and schema metadata, and Paimon consumes the writer-provided IPC schema so nested nullability and physical layouts remain authoritative. Build nested fields through the existing metadata helper and add JSON/VARIANT to that helper. Preserve Arrow storage types, child names, nullability and encoded values; in particular, MAP keys remain non-nullable and LARGEINT retains its string encoding for the full signed 128-bit range. Metadata is carried at every nesting depth and survives Arrow IPC serialization. No session setting is required. Rolling upgrades use an optional result-sink capability flag. Old FEs omit it, so upgraded BEs return the exact legacy metadata layout for both schema fetches and result streams. Upgraded FEs enable the flag and fill missing markers from the planned Doris output types before comparing and advertising schemas. Physical types, names, nullability, dictionaries, conflicting markers and unrelated metadata still participate in strict schema equality. Mixed old/new BE endpoints produce the same advertised metadata in either order. An old BE's stream can still lack the markers advertised by the upgraded FE. PyArrow does not automatically interpret custom field metadata. Applications can use `doris_type=LARGEINT` to convert the corresponding values to integers without guessing from their text. ### Release note Preserve Doris logical type metadata for LARGEINT, IPV4, IPV6, JSON and VARIANT inside Arrow arrays, maps and structs. JSON and VARIANT now also include the marker at the top level. ### Check List (For Author) - Test: - BE ASAN: all 84 tests passed with `bash run-be-ut.sh --run --filter='*MetadataTest*:ArrowSchemaConvertorTest.*:ArrowBlockConvertorTest.*:ArrowSchemaUtilTest.*:DataTypeSerDeArrowTest.*:VParquetWriterTest.*:ArrowFlightTimestampTest.*' -j 48`. Coverage includes metadata compatibility, independent converter policies, table schemas and field IDs, Arrow SerDe, timestamp bindings, and Parquet writer round trips. - The new polymorphic schema-converter test verifies that nested timestamp and LARGEINT metadata policies stay independent across Doris, Flight and legacy Flight instances. - Before the fix, 8 of the 9 new tests failed on missing metadata; the native-type negative control passed. - Coverage includes all five logical types at the root and in ARRAY/STRUCT/MAP, LARGEINT map keys, deep nesting, names/nullability, signed 128-bit extremes, NULL values, and Arrow IPC round trips. - FE: 13 tests passed (5 new compatibility tests and 8 existing Arrow schema helper tests). Tests cover mixed-version schema fetches in both orders, all-old/all-new endpoints, strict rejection of physical and unrelated metadata differences, and result-sink capability negotiation. The pre-fix FE fails the mixed-version and all-old metadata cases. - Legacy BE schema tests cover all five logical types, LARGEINT map keys, and IPC schema round trips. Populated LARGEINT batches also run through the legacy schema path to verify values, NULLs, and IPC remain compatible. - FE Checkstyle, clang-format 16 and `git diff --check` passed. - The optional ADBC integration script has five independent cases for LARGEINT, IPV4, IPV6, JSON and VARIANT. JSON/VARIANT use ARRAY_REPEAT to support nested values in the default mode. - A temporary FE analyzer test checks all five exact sample queries and asserts ARRAY<JSON>/ARRAY<VARIANT> result types. The original combined query fails analysis; the updated queries pass. - Python syntax/import checks passed; all five live ADBC tests were skipped because no cluster URI was provided. No live-cluster result is claimed. - Behavior changed: additional Arrow field metadata only; physical types and values remain compatible. - Documentation: describe nested metadata and client conversion in the Python Flight SQL sample README. --- .../exec/operator/memory_scratch_sink_operator.cpp | 3 +- be/src/exec/operator/result_sink_operator.cpp | 23 +- be/src/exec/operator/result_sink_operator.h | 3 + .../aggregate/aggregate_function_python_udaf.cpp | 4 +- be/src/format/arrow/arrow_block_convertor.cpp | 15 +- be/src/format/arrow/arrow_block_convertor.h | 9 +- be/src/format/arrow/arrow_row_batch.cpp | 147 ++++++---- be/src/format/arrow/arrow_row_batch.h | 71 ++++- .../parquet/parquet_arrow_block_convertor.cpp | 18 +- .../format/parquet/parquet_arrow_block_convertor.h | 20 ++ be/src/format/table/iceberg/arrow_schema_util.cpp | 46 ++-- be/src/format/table/iceberg/arrow_schema_util.h | 20 +- .../iceberg/iceberg_arrow_block_convertor.cpp | 12 +- .../table/paimon/paimon_arrow_block_convertor.cpp | 16 +- .../table/paimon/paimon_arrow_block_convertor.h | 11 + be/src/udf/python/python_udf_meta.cpp | 3 +- .../core/data_type/common_data_type_serder_test.h | 4 +- .../data_type_serde/data_type_serde_arrow_test.cpp | 31 ++- .../format/arrow/arrow_flight_timestamp_test.cpp | 21 +- be/test/format/arrow/arrow_row_batch_test.cpp | 305 +++++++++++++++++++++ .../table/iceberg/arrow_schema_util_test.cpp | 16 +- .../java/org/apache/doris/planner/ResultSink.java | 3 + .../arrowflight/FlightSqlConnectProcessor.java | 9 +- .../service/arrowflight/FlightSqlSchemaHelper.java | 45 +++ .../doris/planner/ResultSinkMetadataTest.java | 38 +++ .../FlightSqlConnectProcessorSchemaTest.java | 167 +++++++++++ gensrc/thrift/DataSinks.thrift | 2 + samples/arrow-flight-sql/python/README.md | 22 ++ .../python/test_nested_type_metadata.py | 108 ++++++++ 29 files changed, 1034 insertions(+), 158 deletions(-) diff --git a/be/src/exec/operator/memory_scratch_sink_operator.cpp b/be/src/exec/operator/memory_scratch_sink_operator.cpp index 7b3bf3eed0b..3520491ebb2 100644 --- a/be/src/exec/operator/memory_scratch_sink_operator.cpp +++ b/be/src/exec/operator/memory_scratch_sink_operator.cpp @@ -103,8 +103,7 @@ 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)); } - DorisArrowBlockConvertor converter(block, state->timezone(), _timezone_obj, - /*datetime_naive=*/true); + ArrowFlightArrowBlockConvertor converter(block, state->timezone(), _timezone_obj); { SCOPED_TIMER(local_state._get_arrow_schema_timer); // After expr executed, use recaculated schema as final schema diff --git a/be/src/exec/operator/result_sink_operator.cpp b/be/src/exec/operator/result_sink_operator.cpp index 24c9f6ad18e..88e21fafd1a 100644 --- a/be/src/exec/operator/result_sink_operator.cpp +++ b/be/src/exec/operator/result_sink_operator.cpp @@ -57,9 +57,8 @@ Status ResultSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& info) } else { std::shared_ptr<arrow::Schema> arrow_schema; if (p._sink_type == TResultSinkType::ARROW_FLIGHT_PROTOCOL) { - RETURN_IF_ERROR(get_arrow_schema_from_expr_ctxs(_output_vexpr_ctxs, &arrow_schema, - state->timezone(), - /*datetime_naive=*/true)); + RETURN_IF_ERROR(p._arrow_schema_convertor->get_arrow_schema_from_expr_ctxs( + _output_vexpr_ctxs, &arrow_schema)); } VLOG_DEBUG << "create sender in INIT with instance id " << fragment_instance_id; RETURN_IF_ERROR(state->exec_env()->result_mgr()->create_sender( @@ -103,6 +102,8 @@ ResultSinkOperatorX::ResultSinkOperatorX(int operator_id, int node_id, _sink_type(!sink.__isset.type || sink.type == TResultSinkType::MYSQL_PROTOCOL ? TResultSinkType::MYSQL_PROTOCOL : sink.type), + _enable_arrow_type_metadata(sink.__isset.enable_arrow_type_metadata && + sink.enable_arrow_type_metadata), _result_sink_buffer_size_rows(_sink_type == TResultSinkType::ARROW_FLIGHT_PROTOCOL ? config::arrow_flight_result_sink_buffer_size_rows : RESULT_SINK_BUFFER_SIZE), @@ -120,12 +121,22 @@ Status ResultSinkOperatorX::prepare(RuntimeState* state) { // Prepare the exprs to run. RETURN_IF_ERROR(VExpr::prepare(_output_vexpr_ctxs, state, _row_desc)); + if (_sink_type == TResultSinkType::ARROW_FLIGHT_PROTOCOL) { + // Capability negotiation selects one immutable policy for schema fetches and batches. + if (_enable_arrow_type_metadata) { + _arrow_schema_convertor = + std::make_shared<ArrowFlightSchemaConvertor>(state->timezone()); + } else { + _arrow_schema_convertor = + std::make_shared<LegacyArrowFlightSchemaConvertor>(state->timezone()); + } + } + if (state->query_options().enable_parallel_result_sink) { std::shared_ptr<arrow::Schema> arrow_schema; if (_sink_type == TResultSinkType::ARROW_FLIGHT_PROTOCOL) { - RETURN_IF_ERROR(get_arrow_schema_from_expr_ctxs(_output_vexpr_ctxs, &arrow_schema, - state->timezone(), - /*datetime_naive=*/true)); + RETURN_IF_ERROR(_arrow_schema_convertor->get_arrow_schema_from_expr_ctxs( + _output_vexpr_ctxs, &arrow_schema)); } VLOG_DEBUG << "create sender in prepare with query id " << state->query_id(); RETURN_IF_ERROR(state->exec_env()->result_mgr()->create_sender( diff --git a/be/src/exec/operator/result_sink_operator.h b/be/src/exec/operator/result_sink_operator.h index 84c86f1127e..173a5b52e82 100644 --- a/be/src/exec/operator/result_sink_operator.h +++ b/be/src/exec/operator/result_sink_operator.h @@ -25,6 +25,7 @@ #include "runtime/result_block_buffer.h" namespace doris { +class DorisArrowSchemaConvertor; #include "common/compile_check_begin.h" class ResultBlockBufferBase; @@ -167,6 +168,8 @@ private: Status _second_phase_fetch_data(RuntimeState* state, Block* final_block); const TResultSinkType::type _sink_type; + const bool _enable_arrow_type_metadata; + std::shared_ptr<const DorisArrowSchemaConvertor> _arrow_schema_convertor; const int _result_sink_buffer_size_rows; // set file options when sink type is FILE std::unique_ptr<ResultFileOptions> _file_opts = nullptr; diff --git a/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp b/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp index 495daee2018..2e2fbd266fb 100644 --- a/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp +++ b/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp @@ -245,10 +245,10 @@ void AggregatePythonUDAF::create(AggregateDataPtr __restrict place) const { std::call_once(_schema_init_flag, [this]() { std::vector<std::shared_ptr<arrow::Field>> fields; - std::string timezone = TimezoneUtils::default_time_zone; + const DorisArrowSchemaConvertor schema_convertor(TimezoneUtils::default_time_zone); for (size_t i = 0; i < argument_types.size(); ++i) { std::shared_ptr<arrow::DataType> arrow_type; - Status st = convert_to_arrow_type(argument_types[i], &arrow_type, timezone); + Status st = schema_convertor.convert_to_arrow_type(argument_types[i], &arrow_type); if (!st.ok()) { throw doris::Exception(ErrorCode::INTERNAL_ERROR, "Failed to convert argument type {} to Arrow type: {}", i, diff --git a/be/src/format/arrow/arrow_block_convertor.cpp b/be/src/format/arrow/arrow_block_convertor.cpp index 6b2ed5a71e8..2148d0934af 100644 --- a/be/src/format/arrow/arrow_block_convertor.cpp +++ b/be/src/format/arrow/arrow_block_convertor.cpp @@ -443,7 +443,8 @@ Status ArrowBlockConvertor::write_plain_arrow_column(const std::shared_ptr<const int64_t start, int64_t end, const cctz::time_zone& ctz) const { std::shared_ptr<arrow::DataType> plain_arrow_type; - RETURN_IF_ERROR(convert_to_arrow_type(type, &plain_arrow_type, ctz.name())); + RETURN_IF_ERROR( + DorisArrowSchemaConvertor(ctz.name()).convert_to_arrow_type(type, &plain_arrow_type)); 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. @@ -501,8 +502,16 @@ 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_IF_ERROR(DorisArrowSchemaConvertor(_header, _timezone_name) + .get_arrow_schema(&_arrow_schema)); + } + return ArrowBlockConvertor::init(); +} + +Status ArrowFlightArrowBlockConvertor::init() { + if (_arrow_schema == nullptr) { + RETURN_IF_ERROR(ArrowFlightSchemaConvertor(_header, _timezone_name) + .get_arrow_schema(&_arrow_schema)); } return ArrowBlockConvertor::init(); } diff --git a/be/src/format/arrow/arrow_block_convertor.h b/be/src/format/arrow/arrow_block_convertor.h index c923507786d..e0b58d76838 100644 --- a/be/src/format/arrow/arrow_block_convertor.h +++ b/be/src/format/arrow/arrow_block_convertor.h @@ -89,11 +89,10 @@ class DorisArrowBlockConvertor : public ArrowBlockConvertor { public: using ArrowBlockConvertor::ArrowBlockConvertor; DorisArrowBlockConvertor(const Block& header, std::string timezone_name, - const cctz::time_zone& timezone, bool datetime_naive = false) + const cctz::time_zone& timezone) : ArrowBlockConvertor(nullptr, timezone), _header(header.clone_empty()), - _timezone_name(std::move(timezone_name)), - _datetime_naive(datetime_naive) {} + _timezone_name(std::move(timezone_name)) {} Status init() override; Status convert_from_arrow(const std::shared_ptr<arrow::RecordBatch>& batch, @@ -106,16 +105,14 @@ protected: arrow::ArrayBuilder* array_builder, int64_t start, int64_t end, const cctz::time_zone& ctz) const override; -private: Block _header; std::string _timezone_name; - bool _datetime_naive = false; }; class ArrowFlightArrowBlockConvertor final : public DorisArrowBlockConvertor { public: using DorisArrowBlockConvertor::DorisArrowBlockConvertor; - + Status init() override; Status convert_to_arrow(const Block& block, arrow::MemoryPool* pool, std::shared_ptr<arrow::RecordBatch>* result, size_t start_row = 0, size_t end_row = 0) const override; diff --git a/be/src/format/arrow/arrow_row_batch.cpp b/be/src/format/arrow/arrow_row_batch.cpp index f291cf33ac3..a958843e519 100644 --- a/be/src/format/arrow/arrow_row_batch.cpp +++ b/be/src/format/arrow/arrow_row_batch.cpp @@ -48,9 +48,8 @@ namespace doris { -Status convert_to_arrow_type(const DataTypePtr& origin_type, - std::shared_ptr<arrow::DataType>* result, const std::string& timezone, - bool datetime_naive) { +Status DorisArrowSchemaConvertor::convert_to_arrow_type( + const DataTypePtr& origin_type, std::shared_ptr<arrow::DataType>* result) const { auto type = get_serialized_type(origin_type); switch (type->get_primitive_type()) { case TYPE_NULL: @@ -107,18 +106,8 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type, } else { time_unit = arrow::TimeUnit::SECOND; } - // Doris DATETIMEV2 represents a wall-clock value without a timezone. Arrow Flight - // exposes it as a timezone-naive timestamp so clients do not interpret it as an instant. - // This option only changes the DATETIMEV2 output schema. TIMESTAMPTZ remains timezone-aware, - // and Arrow-to-Doris conversions are unaffected. - if (type->get_primitive_type() == TYPE_DATETIMEV2 && datetime_naive) { - *result = std::make_shared<arrow::TimestampType>(time_unit); - } else { - // 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<arrow::TimestampType>(time_unit, - timezone == "Z" ? "UTC" : timezone); - } + *result = std::make_shared<arrow::TimestampType>( + time_unit, timestamp_timezone(type->get_primitive_type())); break; } case TYPE_DECIMALV2: @@ -136,20 +125,23 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type, case TYPE_ARRAY: { const auto* type_arr = assert_cast<const DataTypeArray*>(remove_nullable(type).get()); std::shared_ptr<arrow::DataType> item_type; - RETURN_IF_ERROR(convert_to_arrow_type(type_arr->get_nested_type(), &item_type, timezone, - datetime_naive)); - *result = std::make_shared<arrow::ListType>(item_type); + RETURN_IF_ERROR(convert_to_arrow_type(type_arr->get_nested_type(), &item_type)); + // Arrow stores metadata on fields, so implicit child fields lose the Doris logical type. + *result = std::make_shared<arrow::ListType>(make_child_field( + "item", item_type, true, type_arr->get_nested_type()->get_primitive_type())); break; } case TYPE_MAP: { const auto* type_map = assert_cast<const DataTypeMap*>(remove_nullable(type).get()); std::shared_ptr<arrow::DataType> key_type; std::shared_ptr<arrow::DataType> val_type; - RETURN_IF_ERROR(convert_to_arrow_type(type_map->get_key_type(), &key_type, timezone, - datetime_naive)); - RETURN_IF_ERROR(convert_to_arrow_type(type_map->get_value_type(), &val_type, timezone, - datetime_naive)); - *result = std::make_shared<arrow::MapType>(key_type, val_type); + RETURN_IF_ERROR(convert_to_arrow_type(type_map->get_key_type(), &key_type)); + RETURN_IF_ERROR(convert_to_arrow_type(type_map->get_value_type(), &val_type)); + auto key_field = make_child_field("key", key_type, false, + type_map->get_key_type()->get_primitive_type()); + auto value_field = make_child_field("value", val_type, true, + type_map->get_value_type()->get_primitive_type()); + *result = std::make_shared<arrow::MapType>(key_field, value_field); break; } case TYPE_STRUCT: { @@ -157,11 +149,10 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type, std::vector<std::shared_ptr<arrow::Field>> fields; for (size_t i = 0; i < type_struct->get_elements().size(); i++) { std::shared_ptr<arrow::DataType> field_type; - RETURN_IF_ERROR(convert_to_arrow_type(type_struct->get_element(i), &field_type, - timezone, datetime_naive)); - fields.push_back( - std::make_shared<arrow::Field>(type_struct->get_element_name(i), field_type, - type_struct->get_element(i)->is_nullable())); + RETURN_IF_ERROR(convert_to_arrow_type(type_struct->get_element(i), &field_type)); + fields.push_back(make_child_field(type_struct->get_element_name(i), field_type, + type_struct->get_element(i)->is_nullable(), + type_struct->get_element(i)->get_primitive_type())); } *result = std::make_shared<arrow::StructType>(fields); break; @@ -187,56 +178,102 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type, return Status::OK(); } -// Helper function to create an Arrow Field with type metadata if applicable, such as IP types +// Logical types sharing Arrow storage need the same marker at the root and every child field. std::shared_ptr<arrow::Field> create_arrow_field_with_metadata( const std::string& field_name, const std::shared_ptr<arrow::DataType>& arrow_type, bool is_nullable, PrimitiveType primitive_type) { - if (primitive_type == PrimitiveType::TYPE_IPV4) { - auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"IPV4"}); - return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata); - } else if (primitive_type == PrimitiveType::TYPE_IPV6) { - auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"IPV6"}); - return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata); - } else if (primitive_type == PrimitiveType::TYPE_LARGEINT) { - auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"LARGEINT"}); - return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata); - } else { + const char* type_name; + switch (primitive_type) { + case TYPE_IPV4: + type_name = "IPV4"; + break; + case TYPE_IPV6: + type_name = "IPV6"; + break; + case TYPE_LARGEINT: + type_name = "LARGEINT"; + break; + case TYPE_JSONB: + type_name = "JSON"; + break; + case TYPE_VARIANT: + type_name = "VARIANT"; + break; + default: return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable); } + auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {type_name}); + return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata); } -Status get_arrow_schema_from_block(const Block& block, std::shared_ptr<arrow::Schema>* result, - const std::string& timezone, bool datetime_naive) { +Status DorisArrowSchemaConvertor::get_arrow_schema_from_block( + const Block& block, std::shared_ptr<arrow::Schema>* result) const { std::vector<std::shared_ptr<arrow::Field>> fields; for (const auto& type_and_name : block) { std::shared_ptr<arrow::DataType> arrow_type; - RETURN_IF_ERROR( - convert_to_arrow_type(type_and_name.type, &arrow_type, timezone, datetime_naive)); - auto field = create_arrow_field_with_metadata(type_and_name.name, arrow_type, - type_and_name.type->is_nullable(), - type_and_name.type->get_primitive_type()); + RETURN_IF_ERROR(convert_to_arrow_type(type_and_name.type, &arrow_type)); + auto field = make_field(type_and_name.name, arrow_type, type_and_name.type->is_nullable(), + type_and_name.type->get_primitive_type()); fields.push_back(field); } *result = arrow::schema(std::move(fields)); return Status::OK(); } -Status get_arrow_schema_from_expr_ctxs(const VExprContextSPtrs& output_vexpr_ctxs, - std::shared_ptr<arrow::Schema>* result, - const std::string& timezone, bool datetime_naive) { +Status DorisArrowSchemaConvertor::get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const { + return get_arrow_schema_from_block(_header, result); +} + +std::string DorisArrowSchemaConvertor::timestamp_timezone(PrimitiveType) const { + // Arrow clients expect a timezone name rather than the ISO-8601 UTC alias. + return _timezone == "Z" ? "UTC" : _timezone; +} + +std::string ArrowFlightSchemaConvertor::timestamp_timezone(PrimitiveType type) const { + // DATETIMEV2 is wall-clock time; TIMESTAMPTZ must still describe an instant. + return type == TYPE_DATETIMEV2 ? "" : DorisArrowSchemaConvertor::timestamp_timezone(type); +} + +std::shared_ptr<arrow::Field> DorisArrowSchemaConvertor::make_field( + const std::string& name, const std::shared_ptr<arrow::DataType>& type, bool nullable, + PrimitiveType primitive) const { + return create_arrow_field_with_metadata(name, type, nullable, primitive); +} + +std::shared_ptr<arrow::Field> DorisArrowSchemaConvertor::make_child_field( + const std::string& name, const std::shared_ptr<arrow::DataType>& type, bool nullable, + PrimitiveType primitive) const { + return create_arrow_field_with_metadata(name, type, nullable, primitive); +} + +std::shared_ptr<arrow::Field> LegacyArrowFlightSchemaConvertor::make_field( + const std::string& name, const std::shared_ptr<arrow::DataType>& type, bool nullable, + PrimitiveType primitive) const { + if (primitive == TYPE_JSONB || primitive == TYPE_VARIANT) { + return arrow::field(name, type, nullable); + } + return DorisArrowSchemaConvertor::make_field(name, type, nullable, primitive); +} + +std::shared_ptr<arrow::Field> LegacyArrowFlightSchemaConvertor::make_child_field( + const std::string& name, const std::shared_ptr<arrow::DataType>& type, bool nullable, + PrimitiveType) const { + return arrow::field(name, type, nullable); +} + +Status DorisArrowSchemaConvertor::get_arrow_schema_from_expr_ctxs( + const VExprContextSPtrs& output_vexpr_ctxs, std::shared_ptr<arrow::Schema>* result) const { std::vector<std::shared_ptr<arrow::Field>> fields; for (int i = 0; i < output_vexpr_ctxs.size(); i++) { std::shared_ptr<arrow::DataType> arrow_type; auto root_expr = output_vexpr_ctxs.at(i)->root(); - RETURN_IF_ERROR(convert_to_arrow_type(root_expr->data_type(), &arrow_type, timezone, - datetime_naive)); + RETURN_IF_ERROR(convert_to_arrow_type(root_expr->data_type(), &arrow_type)); auto field_name = root_expr->is_slot_ref() && !root_expr->expr_label().empty() ? root_expr->expr_label() : fmt::format("{}_{}", root_expr->data_type()->get_name(), i); - auto field = - create_arrow_field_with_metadata(field_name, arrow_type, root_expr->is_nullable(), - root_expr->data_type()->get_primitive_type()); - fields.push_back(field); + auto field = make_field(field_name, arrow_type, root_expr->is_nullable(), + root_expr->data_type()->get_primitive_type()); + fields.push_back(std::move(field)); } *result = arrow::schema(std::move(fields)); return Status::OK(); diff --git a/be/src/format/arrow/arrow_row_batch.h b/be/src/format/arrow/arrow_row_batch.h index d5ba5cb0ed0..60b192716fe 100644 --- a/be/src/format/arrow/arrow_row_batch.h +++ b/be/src/format/arrow/arrow_row_batch.h @@ -19,6 +19,7 @@ #include <memory> #include <string> +#include <utility> #include "common/status.h" #include "core/block/block.h" @@ -43,23 +44,69 @@ constexpr size_t MAX_ARROW_UTF8 = (1ULL << 31); // 2G class RowDescriptor; -// datetime_naive only controls how Doris DATETIMEV2 is represented in the output Arrow schema. -// When enabled, DATETIMEV2 is mapped to a timestamp without a timezone to preserve its wall-clock -// semantics. TIMESTAMPTZ remains timezone-aware, and Arrow-to-Doris conversions are unaffected. -Status convert_to_arrow_type(const DataTypePtr& type, std::shared_ptr<arrow::DataType>* result, - const std::string& timezone, bool datetime_naive = false); +// Each protocol owns its schema source and conversion rules. Table formats retain their +// authoritative schemas instead of reconstructing field IDs or physical layouts from Doris types. +class ArrowSchemaConvertor { +public: + virtual ~ArrowSchemaConvertor() = default; + virtual Status get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const = 0; +}; + +class DorisArrowSchemaConvertor : public ArrowSchemaConvertor { +public: + explicit DorisArrowSchemaConvertor(std::string timezone) : _timezone(std::move(timezone)) {} + DorisArrowSchemaConvertor(const Block& header, std::string timezone) + : _header(header.clone_empty()), _timezone(std::move(timezone)) {} + + Status get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const override; + Status get_arrow_schema_from_block(const Block& block, + std::shared_ptr<arrow::Schema>* result) const; + Status get_arrow_schema_from_expr_ctxs(const VExprContextSPtrs& output_vexpr_ctxs, + std::shared_ptr<arrow::Schema>* result) const; + Status convert_to_arrow_type(const DataTypePtr& type, + std::shared_ptr<arrow::DataType>* result) const; + +protected: + virtual std::string timestamp_timezone(PrimitiveType type) const; + virtual std::shared_ptr<arrow::Field> make_field(const std::string& name, + const std::shared_ptr<arrow::DataType>& type, + bool nullable, PrimitiveType primitive) const; + virtual std::shared_ptr<arrow::Field> make_child_field( + const std::string& name, const std::shared_ptr<arrow::DataType>& type, bool nullable, + PrimitiveType primitive) const; + +private: + Block _header; + const std::string _timezone; +}; + +class ArrowFlightSchemaConvertor : public DorisArrowSchemaConvertor { +public: + using DorisArrowSchemaConvertor::DorisArrowSchemaConvertor; + +protected: + std::string timestamp_timezone(PrimitiveType type) const override; +}; + +// Old FEs require the pre-capability metadata layout, including metadata-free nested fields. +class LegacyArrowFlightSchemaConvertor final : public ArrowFlightSchemaConvertor { +public: + using ArrowFlightSchemaConvertor::ArrowFlightSchemaConvertor; + +protected: + std::shared_ptr<arrow::Field> make_field(const std::string& name, + const std::shared_ptr<arrow::DataType>& type, + bool nullable, PrimitiveType primitive) const override; + std::shared_ptr<arrow::Field> make_child_field(const std::string& name, + const std::shared_ptr<arrow::DataType>& type, + bool nullable, + PrimitiveType primitive) const override; +}; std::shared_ptr<arrow::Field> create_arrow_field_with_metadata( const std::string& field_name, const std::shared_ptr<arrow::DataType>& arrow_type, bool is_nullable, PrimitiveType primitive_type); -Status get_arrow_schema_from_block(const Block& block, std::shared_ptr<arrow::Schema>* result, - const std::string& timezone, bool datetime_naive = false); - -Status get_arrow_schema_from_expr_ctxs(const VExprContextSPtrs& output_vexpr_ctxs, - std::shared_ptr<arrow::Schema>* result, - const std::string& timezone, bool datetime_naive = false); - Status serialize_record_batch(const arrow::RecordBatch& record_batch, std::string* result); Status serialize_arrow_schema(std::shared_ptr<arrow::Schema>* schema, std::string* result); diff --git a/be/src/format/parquet/parquet_arrow_block_convertor.cpp b/be/src/format/parquet/parquet_arrow_block_convertor.cpp index cd8b933195a..156ee634ac5 100644 --- a/be/src/format/parquet/parquet_arrow_block_convertor.cpp +++ b/be/src/format/parquet/parquet_arrow_block_convertor.cpp @@ -24,7 +24,13 @@ namespace doris { #include "common/compile_check_begin.h" -Status ParquetArrowBlockConvertor::init() { +std::string ParquetArrowSchemaConvertor::timestamp_timezone(PrimitiveType type) const { + return type == TYPE_DATETIMEV2 && !_enable_int96_timestamps + ? "" + : DorisArrowSchemaConvertor::timestamp_timezone(type); +} + +Status ParquetArrowSchemaConvertor::get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const { if (_types.size() != _names.size()) { return Status::InvalidArgument("Parquet column names and types must have the same size"); } @@ -34,14 +40,18 @@ Status ParquetArrowBlockConvertor::init() { // INT96 normalization and schema construction must use the same instance's timezone. for (size_t i = 0; i < _types.size(); ++i) { std::shared_ptr<arrow::DataType> type; - RETURN_IF_ERROR( - convert_to_arrow_type(_types[i], &type, _timezone_name, !_enable_int96_timestamps)); + RETURN_IF_ERROR(convert_to_arrow_type(_types[i], &type)); fields.emplace_back(arrow::field(_names[i], type, _types[i]->is_nullable())); } - _arrow_schema = arrow::schema(std::move(fields)); + *result = arrow::schema(std::move(fields)); return Status::OK(); } +Status ParquetArrowBlockConvertor::init() { + return ParquetArrowSchemaConvertor(_types, _names, _timezone_name, _enable_int96_timestamps) + .get_arrow_schema(&_arrow_schema); +} + Status ParquetArrowBlockConvertor::write_column(const DataTypePtr& type, const DataTypeSerDe& serde, const IColumn& column, const NullMap* null_map, const std::shared_ptr<arrow::Field>& field, diff --git a/be/src/format/parquet/parquet_arrow_block_convertor.h b/be/src/format/parquet/parquet_arrow_block_convertor.h index 24a490a565b..28744817f9c 100644 --- a/be/src/format/parquet/parquet_arrow_block_convertor.h +++ b/be/src/format/parquet/parquet_arrow_block_convertor.h @@ -18,9 +18,29 @@ #pragma once #include "format/arrow/arrow_block_convertor.h" +#include "format/arrow/arrow_row_batch.h" namespace doris { +class ParquetArrowSchemaConvertor final : public DorisArrowSchemaConvertor { +public: + ParquetArrowSchemaConvertor(DataTypes types, std::vector<std::string> names, + std::string timezone, bool enable_int96_timestamps) + : DorisArrowSchemaConvertor(std::move(timezone)), + _types(std::move(types)), + _names(std::move(names)), + _enable_int96_timestamps(enable_int96_timestamps) {} + Status get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const override; + +protected: + std::string timestamp_timezone(PrimitiveType type) const override; + +private: + const DataTypes _types; + const std::vector<std::string> _names; + const bool _enable_int96_timestamps; +}; + class ParquetArrowBlockConvertor : public ArrowBlockConvertor { public: ParquetArrowBlockConvertor(DataTypes types, std::vector<std::string> names, diff --git a/be/src/format/table/iceberg/arrow_schema_util.cpp b/be/src/format/table/iceberg/arrow_schema_util.cpp index 84dc29f594d..848ce74f903 100644 --- a/be/src/format/table/iceberg/arrow_schema_util.cpp +++ b/be/src/format/table/iceberg/arrow_schema_util.cpp @@ -24,24 +24,36 @@ namespace doris::iceberg { #include "common/compile_check_begin.h" -const char* ArrowSchemaUtil::PARQUET_FIELD_ID = "PARQUET:field_id"; -const char* ArrowSchemaUtil::ORIGINAL_TYPE = "originalType"; -const char* ArrowSchemaUtil::MAP_TYPE_VALUE = "mapType"; -const char* ArrowSchemaUtil::UUID_TYPE_VALUE = "uuid"; - -Status ArrowSchemaUtil::convert(const Schema* schema, const std::string& timezone, - std::vector<std::shared_ptr<arrow::Field>>& fields) { - for (const auto& column : schema->columns()) { +const char* IcebergArrowSchemaConvertor::PARQUET_FIELD_ID = "PARQUET:field_id"; +const char* IcebergArrowSchemaConvertor::ORIGINAL_TYPE = "originalType"; +const char* IcebergArrowSchemaConvertor::MAP_TYPE_VALUE = "mapType"; +const char* IcebergArrowSchemaConvertor::UUID_TYPE_VALUE = "uuid"; + +Status IcebergArrowSchemaConvertor::get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const { + // Field IDs, Variant storage and timestamp bindings all come from the pinned table schema. + std::vector<std::shared_ptr<arrow::Field>> fields; + RETURN_IF_ERROR(convert_fields(fields)); + auto schema = arrow::schema(std::move(fields)); + if (!_schema_json.empty()) { + schema = schema->WithMetadata( + arrow::KeyValueMetadata::Make({"iceberg.schema"}, {_schema_json})); + } + *result = std::move(schema); + return Status::OK(); +} + +Status IcebergArrowSchemaConvertor::convert_fields( + std::vector<std::shared_ptr<arrow::Field>>& fields) const { + for (const auto& column : _schema.columns()) { std::shared_ptr<arrow::Field> arrow_field; - RETURN_IF_ERROR(convert_to(column, &arrow_field, timezone)); + RETURN_IF_ERROR(convert_to_arrow_field(column, &arrow_field)); fields.push_back(arrow_field); } return Status::OK(); } -Status ArrowSchemaUtil::convert_to(const iceberg::NestedField& field, - std::shared_ptr<arrow::Field>* arrow_field, - const std::string& timezone) { +Status IcebergArrowSchemaConvertor::convert_to_arrow_field( + const iceberg::NestedField& field, std::shared_ptr<arrow::Field>* arrow_field) const { std::shared_ptr<arrow::DataType> arrow_type; std::unordered_map<std::string, std::string> metadata; metadata[PARQUET_FIELD_ID] = std::to_string(field.field_id()); @@ -73,7 +85,7 @@ Status ArrowSchemaUtil::convert_to(const iceberg::NestedField& field, case iceberg::TypeID::TIMESTAMP: { iceberg::TimestampType* t_type = static_cast<iceberg::TimestampType*>(field.field_type()); - std::string real_tz = t_type->should_adjust_to_utc() ? timezone : ""; + std::string real_tz = t_type->should_adjust_to_utc() ? _timezone : ""; arrow_type = std::make_shared<arrow::TimestampType>(arrow::TimeUnit::MICRO, real_tz); break; } @@ -115,7 +127,7 @@ Status ArrowSchemaUtil::convert_to(const iceberg::NestedField& field, StructType* st = field.field_type()->as_struct_type(); for (const auto& column : st->fields()) { std::shared_ptr<arrow::Field> element_field; - RETURN_IF_ERROR(convert_to(column, &element_field, timezone)); + RETURN_IF_ERROR(convert_to_arrow_field(column, &element_field)); element_fields.push_back(element_field); } arrow_type = arrow::struct_(element_fields); @@ -125,7 +137,7 @@ Status ArrowSchemaUtil::convert_to(const iceberg::NestedField& field, case iceberg::TypeID::LIST: { std::shared_ptr<arrow::Field> item_field; ListType* list_type = field.field_type()->as_list_type(); - RETURN_IF_ERROR(convert_to(list_type->element_field(), &item_field, timezone)); + RETURN_IF_ERROR(convert_to_arrow_field(list_type->element_field(), &item_field)); arrow_type = arrow::list(item_field); break; } @@ -134,8 +146,8 @@ Status ArrowSchemaUtil::convert_to(const iceberg::NestedField& field, std::shared_ptr<arrow::Field> key_field; std::shared_ptr<arrow::Field> value_field; MapType* map_type = field.field_type()->as_map_type(); - RETURN_IF_ERROR(convert_to(map_type->key_field(), &key_field, timezone)); - RETURN_IF_ERROR(convert_to(map_type->value_field(), &value_field, timezone)); + RETURN_IF_ERROR(convert_to_arrow_field(map_type->key_field(), &key_field)); + RETURN_IF_ERROR(convert_to_arrow_field(map_type->value_field(), &value_field)); metadata[ORIGINAL_TYPE] = MAP_TYPE_VALUE; arrow_type = std::make_shared<arrow::MapType>(key_field, value_field); break; diff --git a/be/src/format/table/iceberg/arrow_schema_util.h b/be/src/format/table/iceberg/arrow_schema_util.h index 2942edacbaa..9d0f6d022c7 100644 --- a/be/src/format/table/iceberg/arrow_schema_util.h +++ b/be/src/format/table/iceberg/arrow_schema_util.h @@ -19,15 +19,21 @@ #include <arrow/type.h> +#include "format/arrow/arrow_row_batch.h" #include "format/table/iceberg/schema.h" namespace doris::iceberg { #include "common/compile_check_begin.h" -class ArrowSchemaUtil { +class IcebergArrowSchemaConvertor final : public ArrowSchemaConvertor { public: - static Status convert(const Schema* schema, const std::string& timezone, - std::vector<std::shared_ptr<arrow::Field>>& fields); + IcebergArrowSchemaConvertor(const Schema& schema, std::string timezone, + std::string schema_json = {}) + : _schema(schema), + _timezone(std::move(timezone)), + _schema_json(std::move(schema_json)) {} + Status get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const override; + Status convert_fields(std::vector<std::shared_ptr<arrow::Field>>& fields) const; private: static const char* PARQUET_FIELD_ID; @@ -35,9 +41,11 @@ private: static const char* MAP_TYPE_VALUE; static const char* UUID_TYPE_VALUE; - static Status convert_to(const iceberg::NestedField& field, - std::shared_ptr<arrow::Field>* arrow_field, - const std::string& timezone); + Status convert_to_arrow_field(const iceberg::NestedField& field, + std::shared_ptr<arrow::Field>* arrow_field) const; + const Schema& _schema; + const std::string _timezone; + const std::string _schema_json; }; #include "common/compile_check_end.h" diff --git a/be/src/format/table/iceberg/iceberg_arrow_block_convertor.cpp b/be/src/format/table/iceberg/iceberg_arrow_block_convertor.cpp index f5260b76e6d..0cbcebd2ecb 100644 --- a/be/src/format/table/iceberg/iceberg_arrow_block_convertor.cpp +++ b/be/src/format/table/iceberg/iceberg_arrow_block_convertor.cpp @@ -30,16 +30,8 @@ Status IcebergArrowBlockConvertor::init() { if (_schema == nullptr) { return ArrowBlockConvertor::init(); } - // Field IDs, Variant storage and timestamp bindings must share the same target schema. - std::vector<std::shared_ptr<arrow::Field>> 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(); + return IcebergArrowSchemaConvertor(*_schema, _timezone_name, _schema_json) + .get_arrow_schema(&_arrow_schema); } Status IcebergArrowBlockConvertor::write_column(const std::shared_ptr<const IDataType>& type, diff --git a/be/src/format/table/paimon/paimon_arrow_block_convertor.cpp b/be/src/format/table/paimon/paimon_arrow_block_convertor.cpp index 9404d622938..3baa787cd40 100644 --- a/be/src/format/table/paimon/paimon_arrow_block_convertor.cpp +++ b/be/src/format/table/paimon/paimon_arrow_block_convertor.cpp @@ -25,10 +25,7 @@ namespace doris::paimon { #include "common/compile_check_begin.h" -Status PaimonArrowBlockConvertor::init() { - if (_arrow_schema != nullptr) { - return Status::OK(); - } +Status PaimonArrowSchemaConvertor::get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const { // 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::io::BufferReader>( @@ -38,7 +35,16 @@ Status PaimonArrowBlockConvertor::init() { return Status::InvalidArgument("Failed to deserialize Paimon Arrow schema: {}", reader.status().ToString()); } - _arrow_schema = reader.ValueOrDie()->schema(); + *result = reader.ValueOrDie()->schema(); + return Status::OK(); +} + +Status PaimonArrowBlockConvertor::init() { + if (_arrow_schema != nullptr) { + return Status::OK(); + } + RETURN_IF_ERROR( + PaimonArrowSchemaConvertor(_serialized_schema).get_arrow_schema(&_arrow_schema)); _serialized_schema.clear(); return Status::OK(); } diff --git a/be/src/format/table/paimon/paimon_arrow_block_convertor.h b/be/src/format/table/paimon/paimon_arrow_block_convertor.h index caa6501e0a6..8dd159dca3f 100644 --- a/be/src/format/table/paimon/paimon_arrow_block_convertor.h +++ b/be/src/format/table/paimon/paimon_arrow_block_convertor.h @@ -18,9 +18,20 @@ #pragma once #include "format/arrow/arrow_block_convertor.h" +#include "format/arrow/arrow_row_batch.h" namespace doris::paimon { +class PaimonArrowSchemaConvertor final : public ArrowSchemaConvertor { +public: + explicit PaimonArrowSchemaConvertor(std::string serialized_schema) + : _serialized_schema(std::move(serialized_schema)) {} + Status get_arrow_schema(std::shared_ptr<arrow::Schema>* result) const override; + +private: + const std::string _serialized_schema; +}; + // 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. diff --git a/be/src/udf/python/python_udf_meta.cpp b/be/src/udf/python/python_udf_meta.cpp index cbf885094a2..e0cd98ba01c 100644 --- a/be/src/udf/python/python_udf_meta.cpp +++ b/be/src/udf/python/python_udf_meta.cpp @@ -34,9 +34,10 @@ namespace doris { Status PythonUDFMeta::convert_types_to_schema(const DataTypes& types, const std::string& timezone, std::shared_ptr<arrow::Schema>* schema) { arrow::SchemaBuilder builder; + const DorisArrowSchemaConvertor schema_convertor(timezone); for (size_t i = 0; i < types.size(); ++i) { std::shared_ptr<arrow::DataType> arrow_type; - RETURN_IF_ERROR(convert_to_arrow_type(types[i], &arrow_type, timezone)); + RETURN_IF_ERROR(schema_convertor.convert_to_arrow_type(types[i], &arrow_type)); std::shared_ptr<arrow::Field> field = std::make_shared<arrow::Field>( "arg" + std::to_string(i), arrow_type, types[i]->is_nullable()); RETURN_DORIS_STATUS_IF_ERROR(builder.AddField(field)); diff --git a/be/test/core/data_type/common_data_type_serder_test.h b/be/test/core/data_type/common_data_type_serder_test.h index ad8667f1454..1760e0d6c39 100644 --- a/be/test/core/data_type/common_data_type_serder_test.h +++ b/be/test/core/data_type/common_data_type_serder_test.h @@ -362,7 +362,9 @@ public: static std::shared_ptr<arrow::RecordBatch> serialize_arrow( const std::shared_ptr<Block>& block) { std::shared_ptr<arrow::Schema> block_arrow_schema; - EXPECT_EQ(get_arrow_schema_from_block(*block, &block_arrow_schema, "UTC"), Status::OK()); + EXPECT_EQ(DorisArrowSchemaConvertor("UTC").get_arrow_schema_from_block(*block, + &block_arrow_schema), + Status::OK()); std::cout << "schema: " << block_arrow_schema->ToString(true) << std::endl; // convert block to arrow std::shared_ptr<arrow::RecordBatch> result; 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 fed7e4c7be3..d83b2ff8f8e 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 @@ -550,7 +550,8 @@ void block_converter_test(std::vector<PrimitiveType> cols, int row_num, bool is_ std::shared_ptr<arrow::RecordBatch> record_batch; std::shared_ptr<arrow::Schema> schema; Status status = Status::OK(); - status = get_arrow_schema_from_block(*source_block, &schema, TimezoneUtils::default_time_zone); + status = DorisArrowSchemaConvertor(TimezoneUtils::default_time_zone) + .get_arrow_schema_from_block(*source_block, &schema); ASSERT_TRUE(status.ok() && schema); cctz::time_zone default_timezone; ASSERT_TRUE( @@ -587,7 +588,8 @@ TEST(DataTypeSerDeArrowTest, DataTypeCollectionSerDeTest) { TEST(DataTypeSerDeArrowTest, ArrowBlockConvertorReusesBothDirectionsAndValidatesSlices) { auto source = create_test_block({TYPE_INT, TYPE_STRING, TYPE_VARBINARY}, 4, true); std::shared_ptr<arrow::Schema> schema; - ASSERT_TRUE(get_arrow_schema_from_block(*source, &schema, "UTC").ok()); + ASSERT_TRUE( + DorisArrowSchemaConvertor("UTC").get_arrow_schema_from_block(*source, &schema).ok()); ArrowFlightArrowBlockConvertor convertor(schema, cctz::utc_time_zone()); std::shared_ptr<arrow::RecordBatch> batch; ASSERT_TRUE(convertor.convert_to_arrow(*source, arrow::default_memory_pool(), &batch).ok()); @@ -607,7 +609,7 @@ void expect_target_converter_matches_plain(const std::vector<PrimitiveType>& typ ArrowConvertorFactory target_converter) { auto block = create_test_block(types, 4, false); std::shared_ptr<arrow::Schema> schema; - ASSERT_TRUE(get_arrow_schema_from_block(*block, &schema, "UTC").ok()); + ASSERT_TRUE(DorisArrowSchemaConvertor("UTC").get_arrow_schema_from_block(*block, &schema).ok()); std::shared_ptr<arrow::RecordBatch> plain_batch; ASSERT_TRUE( @@ -651,7 +653,7 @@ TEST(DataTypeSerDeArrowTest, PaimonCommonScalarTypesUseDeclaredConverter) { TEST(DataTypeSerDeArrowTest, PlainArrowWritesAggregateStateBinaryTypes) { auto block = create_test_block({TYPE_HLL, TYPE_BITMAP, TYPE_QUANTILE_STATE}, 3, false); std::shared_ptr<arrow::Schema> schema; - ASSERT_TRUE(get_arrow_schema_from_block(*block, &schema, "UTC").ok()); + ASSERT_TRUE(DorisArrowSchemaConvertor("UTC").get_arrow_schema_from_block(*block, &schema).ok()); std::shared_ptr<arrow::RecordBatch> batch; Status status = convert_to_arrow_batch_for_test( *block, schema, arrow::default_memory_pool(), &batch, cctz::utc_time_zone(), 0, @@ -668,7 +670,7 @@ TEST(DataTypeSerDeArrowTest, PlainArrowWritesAggregateStateBinaryTypes) { TEST(DataTypeSerDeArrowTest, PlainArrowWritesTimeV2) { auto block = create_test_block({TYPE_TIMEV2}, 3, false); std::shared_ptr<arrow::Schema> schema; - ASSERT_TRUE(get_arrow_schema_from_block(*block, &schema, "UTC").ok()); + ASSERT_TRUE(DorisArrowSchemaConvertor("UTC").get_arrow_schema_from_block(*block, &schema).ok()); ASSERT_EQ(arrow::Type::DOUBLE, schema->field(0)->type()->id()); std::shared_ptr<arrow::RecordBatch> batch; @@ -1428,27 +1430,31 @@ TEST(DataTypeSerDeArrowTest, ConvertDateTimeV2ToNaiveArrowType) { const auto datetime_type = std::make_shared<DataTypeDateTimeV2>(6); std::shared_ptr<arrow::DataType> arrow_type; - auto status = convert_to_arrow_type(datetime_type, &arrow_type, "Asia/Shanghai"); + auto status = DorisArrowSchemaConvertor("Asia/Shanghai") + .convert_to_arrow_type(datetime_type, &arrow_type); ASSERT_TRUE(status.ok()) << status; auto timestamp_type = std::static_pointer_cast<arrow::TimestampType>(arrow_type); EXPECT_EQ(arrow::TimeUnit::MICRO, timestamp_type->unit()); EXPECT_EQ("Asia/Shanghai", timestamp_type->timezone()); - status = convert_to_arrow_type(datetime_type, &arrow_type, "Asia/Shanghai", true); + status = ArrowFlightSchemaConvertor("Asia/Shanghai") + .convert_to_arrow_type(datetime_type, &arrow_type); ASSERT_TRUE(status.ok()) << status; timestamp_type = std::static_pointer_cast<arrow::TimestampType>(arrow_type); EXPECT_EQ(arrow::TimeUnit::MICRO, timestamp_type->unit()); EXPECT_TRUE(timestamp_type->timezone().empty()); const auto timestamptz_type = std::make_shared<DataTypeTimeStampTz>(6); - status = convert_to_arrow_type(timestamptz_type, &arrow_type, "Asia/Shanghai", true); + status = ArrowFlightSchemaConvertor("Asia/Shanghai") + .convert_to_arrow_type(timestamptz_type, &arrow_type); ASSERT_TRUE(status.ok()) << status; timestamp_type = std::static_pointer_cast<arrow::TimestampType>(arrow_type); EXPECT_EQ(arrow::TimeUnit::MICRO, timestamp_type->unit()); EXPECT_EQ("Asia/Shanghai", timestamp_type->timezone()); const auto array_type = std::make_shared<DataTypeArray>(datetime_type); - status = convert_to_arrow_type(array_type, &arrow_type, "Asia/Shanghai", true); + status = ArrowFlightSchemaConvertor("Asia/Shanghai") + .convert_to_arrow_type(array_type, &arrow_type); ASSERT_TRUE(status.ok()) << status; const auto list_type = std::static_pointer_cast<arrow::ListType>(arrow_type); timestamp_type = std::static_pointer_cast<arrow::TimestampType>(list_type->value_type()); @@ -1459,7 +1465,8 @@ TEST(DataTypeSerDeArrowTest, CanonicalizeUtcTimezoneForArrow) { const auto timestamptz_type = std::make_shared<DataTypeTimeStampTz>(6); std::shared_ptr<arrow::DataType> arrow_type; - auto status = convert_to_arrow_type(timestamptz_type, &arrow_type, "Z", true); + auto status = + ArrowFlightSchemaConvertor("Z").convert_to_arrow_type(timestamptz_type, &arrow_type); ASSERT_TRUE(status.ok()) << status; const auto timestamp_type = std::static_pointer_cast<arrow::TimestampType>(arrow_type); EXPECT_EQ("UTC", timestamp_type->timezone()); @@ -1544,7 +1551,9 @@ TEST(DataTypeSerDeArrowTest, NestedDateTimeV2PlainArrowAcceptsNaiveSchema) { block.insert(ColumnWithTypeAndName(struct_column->get_ptr(), struct_type, "event_struct")); std::shared_ptr<arrow::Schema> naive_schema; - ASSERT_TRUE(get_arrow_schema_from_block(block, &naive_schema, "Asia/Shanghai", true).ok()); + ASSERT_TRUE(ArrowFlightSchemaConvertor("Asia/Shanghai") + .get_arrow_schema_from_block(block, &naive_schema) + .ok()); std::shared_ptr<arrow::RecordBatch> naive_batch; const auto status = convert_to_arrow_batch_for_test( block, naive_schema, arrow::default_memory_pool(), &naive_batch, diff --git a/be/test/format/arrow/arrow_flight_timestamp_test.cpp b/be/test/format/arrow/arrow_flight_timestamp_test.cpp index ebcafc04a91..1d278f75503 100644 --- a/be/test/format/arrow/arrow_flight_timestamp_test.cpp +++ b/be/test/format/arrow/arrow_flight_timestamp_test.cpp @@ -26,6 +26,7 @@ #include "core/data_type/data_type_nullable.h" #include "core/data_type/data_type_struct.h" #include "format/arrow/arrow_block_convertor.h" +#include "format/arrow/arrow_row_batch.h" #include "util/timezone_utils.h" namespace doris { @@ -59,7 +60,7 @@ TEST_F(ArrowFlightTimestampTest, RejectsOutOfRangeInEveryUnitWithoutPublishingBa make_datetime(10000, 1, 1, 0, 0, 0, 0)}) { SCOPED_TRACE(scale); auto block = timestamp_block(scale, {value}); - ArrowFlightArrowBlockConvertor flight(block, "UTC", cctz::utc_time_zone(), true); + ArrowFlightArrowBlockConvertor flight(block, "UTC", cctz::utc_time_zone()); ASSERT_TRUE(flight.init().ok()); const ArrowBlockConvertor& converter = flight; std::shared_ptr<arrow::RecordBatch> batch; @@ -72,7 +73,7 @@ TEST_F(ArrowFlightTimestampTest, RejectsOutOfRangeInEveryUnitWithoutPublishingBa EXPECT_EQ(nullptr, batch); // Other Arrow consumers retain their existing date semantics. - DorisArrowBlockConvertor ordinary(block, "UTC", cctz::utc_time_zone(), true); + DorisArrowBlockConvertor ordinary(flight.arrow_schema(), cctz::utc_time_zone()); ASSERT_TRUE(ordinary.init().ok()); ASSERT_TRUE( ordinary.convert_to_arrow(block, arrow::default_memory_pool(), &batch).ok()); @@ -88,7 +89,7 @@ TEST_F(ArrowFlightTimestampTest, PreservesCalendarBoundariesAndPreEpochFractions auto block = timestamp_block(scale, {make_datetime(1, 1, 1, 0, 0, 0, 0), make_datetime(9999, 12, 31, 23, 59, 59, fraction), make_datetime(1969, 12, 31, 23, 59, 59, fraction)}); - ArrowFlightArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone(), true); + ArrowFlightArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone()); ASSERT_TRUE(converter.init().ok()); std::shared_ptr<arrow::RecordBatch> batch; ASSERT_TRUE(converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch).ok()); @@ -101,7 +102,7 @@ TEST_F(ArrowFlightTimestampTest, PreservesCalendarBoundariesAndPreEpochFractions TEST_F(ArrowFlightTimestampTest, ChecksSlicesAndSubsequentBatches) { auto block = timestamp_block(6, {make_datetime(2024, 1, 1, 0, 0, 0, 0), DateTime {}}); - ArrowFlightArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone(), true); + ArrowFlightArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone()); ASSERT_TRUE(converter.init().ok()); std::shared_ptr<arrow::RecordBatch> batch; ASSERT_TRUE(converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch, 0, 1).ok()); @@ -138,7 +139,7 @@ TEST_F(ArrowFlightTimestampTest, RejectsNestedTimestampValuesAndMapKeys) { column->insert_default(); column->insert(fields[i]); Block block {{std::move(column), types[i], "nested"}}; - ArrowFlightArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone(), true); + ArrowFlightArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone()); ASSERT_TRUE(converter.init().ok()); std::shared_ptr<arrow::RecordBatch> batch; const auto status = @@ -161,7 +162,11 @@ TEST_F(ArrowFlightTimestampTest, ChecksBothUtcAndZonedCalendarBounds) { cctz::time_zone timezone; ASSERT_TRUE(TimezoneUtils::find_cctz_time_zone(zone, timezone)); auto block = timestamp_block(6, {value}); - ArrowFlightArrowBlockConvertor converter(block, zone, timezone); + // Flight-generated DATETIMEV2 schemas are naive; keep an explicit zoned schema + // here so this test still checks both UTC and local calendar bounds. + std::shared_ptr<arrow::Schema> schema; + ASSERT_TRUE(DorisArrowSchemaConvertor(block, zone).get_arrow_schema(&schema).ok()); + ArrowFlightArrowBlockConvertor converter(schema, timezone); ASSERT_TRUE(converter.init().ok()); std::shared_ptr<arrow::RecordBatch> batch; const auto status = converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch); @@ -179,7 +184,7 @@ TEST_F(ArrowFlightTimestampTest, NaiveBoundsDoNotDependOnSessionTimezone) { ASSERT_TRUE(TimezoneUtils::find_cctz_time_zone("+08:00", timezone)); auto block = timestamp_block(6, {make_datetime(1, 1, 1, 0, 0, 0, 0), make_datetime(9999, 12, 31, 23, 59, 59, 999999)}); - ArrowFlightArrowBlockConvertor converter(block, "+08:00", timezone, true); + ArrowFlightArrowBlockConvertor converter(block, "+08:00", timezone); ASSERT_TRUE(converter.init().ok()); std::shared_ptr<arrow::RecordBatch> batch; ASSERT_TRUE(converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch).ok()); @@ -207,7 +212,7 @@ TEST_F(ArrowFlightTimestampTest, IgnoresTimestampsMaskedByNullParents) { nulls->insert_value(1); Block block {{ColumnNullable::create(std::move(data), std::move(nulls)), make_nullable(types[i]), "masked"}}; - ArrowFlightArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone(), true); + ArrowFlightArrowBlockConvertor converter(block, "UTC", cctz::utc_time_zone()); ASSERT_TRUE(converter.init().ok()); std::shared_ptr<arrow::RecordBatch> batch; const auto status = converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch); diff --git a/be/test/format/arrow/arrow_row_batch_test.cpp b/be/test/format/arrow/arrow_row_batch_test.cpp new file mode 100644 index 00000000000..8b720c675d8 --- /dev/null +++ b/be/test/format/arrow/arrow_row_batch_test.cpp @@ -0,0 +1,305 @@ +// 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_row_batch.h" + +#include <arrow/api.h> +#include <arrow/io/api.h> +#include <arrow/ipc/api.h> +#include <gtest/gtest.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_string.h" +#include "core/data_type/data_type_struct.h" +#include "exprs/vexpr_context.h" +#include "exprs/vslot_ref.h" +#include "format/arrow/arrow_block_convertor.h" + +namespace doris { +namespace { + +void expect_logical_type(const std::shared_ptr<arrow::Field>& field, const std::string& name) { + ASSERT_NE(nullptr, field->metadata()) << field->ToString(); + auto value = field->metadata()->Get("doris_type"); + ASSERT_TRUE(value.ok()) << value.status(); + EXPECT_EQ(name, *value); +} + +void expect_legacy_batch(const Block& block, const std::shared_ptr<arrow::RecordBatch>& extended) { + VExprContextSPtrs expressions; + for (size_t i = 0; i < block.columns(); ++i) { + const auto& column = block.get_by_position(i); + expressions.push_back(VExprContext::create_shared( + VSlotRef::create_shared(i, i, -1, column.type, column.name))); + } + std::shared_ptr<arrow::Schema> schema; + ASSERT_TRUE(LegacyArrowFlightSchemaConvertor("UTC") + .get_arrow_schema_from_expr_ctxs(expressions, &schema) + .ok()); + // Synthetic slot references have no expression labels; compare batches using the same names. + schema = schema->WithNames(extended->schema()->field_names()).ValueOrDie(); + std::shared_ptr<arrow::RecordBatch> batch; + auto status = ArrowFlightArrowBlockConvertor(schema, cctz::utc_time_zone()) + .convert_to_arrow(block, arrow::default_memory_pool(), &batch); + ASSERT_TRUE(status.ok()) << status; + ASSERT_TRUE(batch->ValidateFull().ok()); + EXPECT_TRUE(extended->Equals(*batch, false)); + EXPECT_FALSE(extended->schema()->Equals(*schema, true)); + std::string serialized; + ASSERT_TRUE(serialize_record_batch(*batch, &serialized).ok()); + auto source = std::make_shared<arrow::io::BufferReader>(arrow::Buffer::FromString(serialized)); + auto reader = arrow::ipc::RecordBatchStreamReader::Open(source).ValueOrDie(); + ASSERT_TRUE(reader->ReadNext(&batch).ok()); + EXPECT_TRUE(schema->Equals(*batch->schema(), true)); + EXPECT_TRUE(extended->Equals(*batch, false)); +} + +TEST(ArrowSchemaConvertorTest, ProtocolsKeepNestedTimestampAndMetadataPoliciesSeparate) { + const auto datetime = + DataTypeFactory::instance().create_data_type(TYPE_DATETIMEV2, false, 0, 6); + const auto largeint = std::make_shared<DataTypeInt128>(); + const auto nested = std::make_shared<DataTypeStruct>(DataTypes {datetime, largeint}, + Strings {"time", "number"}); + Block block; + block.insert({nested->create_column(), nested, "value"}); + std::vector<std::unique_ptr<ArrowSchemaConvertor>> convertors; + convertors.push_back(std::make_unique<DorisArrowSchemaConvertor>(block, "Asia/Shanghai")); + convertors.push_back(std::make_unique<ArrowFlightSchemaConvertor>(block, "Asia/Shanghai")); + convertors.push_back( + std::make_unique<LegacyArrowFlightSchemaConvertor>(block, "Asia/Shanghai")); + for (size_t i = 0; i < convertors.size(); ++i) { + std::shared_ptr<arrow::Schema> schema; + ASSERT_TRUE(convertors[i]->get_arrow_schema(&schema).ok()); + const auto& fields = schema->field(0)->type()->fields(); + const auto& timestamp = static_cast<const arrow::TimestampType&>(*fields[0]->type()); + EXPECT_EQ(i == 0 ? "Asia/Shanghai" : "", timestamp.timezone()); + if (i == 2) { + EXPECT_EQ(nullptr, fields[1]->metadata()); + } else { + expect_logical_type(fields[1], "LARGEINT"); + } + } +} + +class ArrowLogicalTypeMetadataTest + : public testing::TestWithParam<std::pair<PrimitiveType, const char*>> {}; + +TEST_P(ArrowLogicalTypeMetadataTest, PreservesTopLevelAndNestedFields) { + const auto [primitive, name] = GetParam(); + auto type = make_nullable(DataTypeFactory::instance().create_data_type(primitive, false)); + auto string_type = std::make_shared<DataTypeString>(); + DataTypes types {type, std::make_shared<DataTypeArray>(type), + std::make_shared<DataTypeStruct>(DataTypes {type, string_type}, + Strings {"typed", "text"}), + std::make_shared<DataTypeMap>(string_type, type)}; + Block block; + for (size_t i = 0; i < types.size(); ++i) { + block.insert({types[i]->create_column(), types[i], std::to_string(i)}); + } + std::shared_ptr<arrow::Schema> schema; + ASSERT_TRUE(ArrowFlightSchemaConvertor("UTC").get_arrow_schema_from_block(block, &schema).ok()); + expect_logical_type(schema->field(0), name); + expect_logical_type(schema->field(1)->type()->field(0), name); + expect_logical_type(schema->field(2)->type()->field(0), name); + EXPECT_EQ(nullptr, schema->field(2)->type()->field(1)->metadata()); + const auto& map = static_cast<const arrow::MapType&>(*schema->field(3)->type()); + expect_logical_type(map.item_field(), name); + EXPECT_FALSE(map.key_field()->nullable()); + EXPECT_TRUE(map.item_field()->nullable()); + EXPECT_EQ(nullptr, map.key_field()->metadata()); +} + +TEST_P(ArrowLogicalTypeMetadataTest, OldFeReceivesLegacySchema) { + const auto [primitive, name] = GetParam(); + auto type = make_nullable(DataTypeFactory::instance().create_data_type(primitive, false)); + auto string_type = std::make_shared<DataTypeString>(); + DataTypes types { + type, std::make_shared<DataTypeArray>(type), + std::make_shared<DataTypeStruct>(DataTypes {type}, Strings {"typed"}), + std::make_shared<DataTypeMap>(primitive == TYPE_LARGEINT ? type : string_type, type)}; + VExprContextSPtrs expressions; + for (size_t i = 0; i < types.size(); ++i) { + expressions.push_back(VExprContext::create_shared( + VSlotRef::create_shared(i, i, -1, types[i], std::to_string(i)))); + } + std::shared_ptr<arrow::Schema> legacy; + std::shared_ptr<arrow::Schema> extended; + ASSERT_TRUE(LegacyArrowFlightSchemaConvertor("UTC") + .get_arrow_schema_from_expr_ctxs(expressions, &legacy) + .ok()); + ASSERT_TRUE(ArrowFlightSchemaConvertor("UTC") + .get_arrow_schema_from_expr_ctxs(expressions, &extended) + .ok()); + std::shared_ptr<arrow::DataType> storage; + ASSERT_TRUE(ArrowFlightSchemaConvertor("UTC").convert_to_arrow_type(type, &storage).ok()); + auto scalar = arrow::field(legacy->field(0)->name(), storage, type->is_nullable()); + if (primitive != TYPE_JSONB && primitive != TYPE_VARIANT) { + scalar = scalar->WithMetadata(arrow::key_value_metadata({"doris_type"}, {name})); + } + auto expected = arrow::schema( + {scalar, + arrow::field(legacy->field(1)->name(), arrow::list(storage), types[1]->is_nullable()), + arrow::field(legacy->field(2)->name(), + arrow::struct_({arrow::field("typed", storage, true)}), + types[2]->is_nullable()), + arrow::field(legacy->field(3)->name(), arrow::map(arrow::utf8(), storage), + types[3]->is_nullable())}); + EXPECT_TRUE(expected->Equals(*legacy, true)); + EXPECT_TRUE(legacy->Equals(*extended, false)); + EXPECT_FALSE(legacy->Equals(*extended, true)); + expect_logical_type(extended->field(0), name); + expect_logical_type(extended->field(1)->type()->field(0), name); + std::string serialized; + ASSERT_TRUE(serialize_arrow_schema(&legacy, &serialized).ok()); + auto source = std::make_shared<arrow::io::BufferReader>(arrow::Buffer::FromString(serialized)); + auto reader = arrow::ipc::RecordBatchStreamReader::Open(source).ValueOrDie(); + EXPECT_TRUE(expected->Equals(*reader->schema(), true)); +} + +INSTANTIATE_TEST_SUITE_P(LogicalTypes, ArrowLogicalTypeMetadataTest, + testing::Values(std::make_pair(TYPE_LARGEINT, "LARGEINT"), + std::make_pair(TYPE_IPV4, "IPV4"), + std::make_pair(TYPE_IPV6, "IPV6"), + std::make_pair(TYPE_JSONB, "JSON"), + std::make_pair(TYPE_VARIANT, "VARIANT"))); + +TEST(ArrowRowBatchMetadataTest, PreservesMapKeysAndDeepNestingThroughIpc) { + auto integer = std::make_shared<DataTypeInt128>(); + auto structure = std::make_shared<DataTypeStruct>(DataTypes {integer}, Strings {"number"}); + auto array = std::make_shared<DataTypeArray>(structure); + auto type = std::make_shared<DataTypeMap>(integer, array); + Block block; + block.insert({type->create_column(), type, "m"}); + std::shared_ptr<arrow::Schema> schema; + ASSERT_TRUE(DorisArrowSchemaConvertor("UTC").get_arrow_schema_from_block(block, &schema).ok()); + std::string serialized; + ASSERT_TRUE(serialize_arrow_schema(&schema, &serialized).ok()); + auto source = std::make_shared<arrow::io::BufferReader>(arrow::Buffer::FromString(serialized)); + auto reader = arrow::ipc::RecordBatchStreamReader::Open(source).ValueOrDie(); + EXPECT_TRUE(schema->Equals(*reader->schema(), true)); + const auto& map = static_cast<const arrow::MapType&>(*reader->schema()->field(0)->type()); + expect_logical_type(map.key_field(), "LARGEINT"); + EXPECT_EQ("key", map.key_field()->name()); + EXPECT_FALSE(map.key_field()->nullable()); + EXPECT_EQ("value", map.item_field()->name()); + auto item = map.item_type()->field(0); + EXPECT_EQ("item", item->name()); + EXPECT_TRUE(item->nullable()); + expect_logical_type(item->type()->field(0), "LARGEINT"); + EXPECT_FALSE(item->type()->field(0)->nullable()); + auto physical = + arrow::map(arrow::utf8(), + arrow::list(arrow::struct_({arrow::field("number", arrow::utf8(), false)}))); + EXPECT_TRUE(map.Equals(physical)); + EXPECT_FALSE(map.Equals(physical, true)); +} + +TEST(ArrowRowBatchMetadataTest, KeepsNativeTypesUnannotated) { + auto type = std::make_shared<DataTypeStruct>( + DataTypes {std::make_shared<DataTypeInt32>(), std::make_shared<DataTypeString>()}, + Strings {"number", "text"}); + std::shared_ptr<arrow::DataType> arrow_type; + ASSERT_TRUE(DorisArrowSchemaConvertor("UTC").convert_to_arrow_type(type, &arrow_type).ok()); + for (const auto& field : arrow_type->fields()) { + EXPECT_EQ(nullptr, field->metadata()); + EXPECT_FALSE(field->nullable()); + } +} + +TEST(ArrowRowBatchMetadataTest, PreservesLargeintExtremesAndNullsInRecordBatches) { + auto integer = make_nullable(std::make_shared<DataTypeInt128>()); + auto array_type = std::make_shared<DataTypeArray>(integer); + auto column = array_type->create_column(); + column->insert(Field::create_field<TYPE_ARRAY>( + Array {Field::create_field<TYPE_LARGEINT>(MAX_INT128), + Field::create_field<TYPE_LARGEINT>(MIN_INT128), Field()})); + column->insert(Field::create_field<TYPE_ARRAY>(Array {})); + Block block; + block.insert({std::move(column), array_type, "numbers"}); + std::shared_ptr<arrow::Schema> schema; + ASSERT_TRUE(ArrowFlightSchemaConvertor("UTC").get_arrow_schema_from_block(block, &schema).ok()); + ArrowFlightArrowBlockConvertor converter(schema, cctz::utc_time_zone()); + std::shared_ptr<arrow::RecordBatch> batch; + auto status = converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch); + ASSERT_TRUE(status.ok()) << status; + ASSERT_TRUE(batch->ValidateFull().ok()); + std::string serialized; + ASSERT_TRUE(serialize_record_batch(*batch, &serialized).ok()); + auto source = std::make_shared<arrow::io::BufferReader>(arrow::Buffer::FromString(serialized)); + auto reader = arrow::ipc::RecordBatchStreamReader::Open(source).ValueOrDie(); + ASSERT_TRUE(reader->ReadNext(&batch).ok()); + EXPECT_TRUE(batch->schema()->Equals(*schema, true)); + expect_logical_type(batch->schema()->field(0)->type()->field(0), "LARGEINT"); + const auto& lists = static_cast<const arrow::ListArray&>(*batch->column(0)); + EXPECT_EQ(0, lists.value_length(1)); + const auto& values = static_cast<const arrow::StringArray&>(*lists.values()); + EXPECT_EQ("170141183460469231731687303715884105727", values.GetString(0)); + EXPECT_EQ("-170141183460469231731687303715884105728", values.GetString(1)); + EXPECT_TRUE(values.IsNull(2)); + expect_legacy_batch(block, batch); +} + +TEST(ArrowRowBatchMetadataTest, PreservesNestedMapMetadataAndValuesInRecordBatches) { + auto integer = make_nullable(std::make_shared<DataTypeInt128>()); + auto map_type = make_nullable(std::make_shared<DataTypeMap>(integer, integer)); + auto type = make_nullable(std::make_shared<DataTypeStruct>( + DataTypes {map_type, std::make_shared<DataTypeString>()}, Strings {"mapping", "text"})); + Map map; + map.push_back(Field::create_field<TYPE_ARRAY>( + Array {Field::create_field<TYPE_LARGEINT>(MAX_INT128)})); + map.push_back(Field::create_field<TYPE_ARRAY>( + Array {Field::create_field<TYPE_LARGEINT>(MIN_INT128)})); + auto column = type->create_column(); + column->insert(Field::create_field<TYPE_STRUCT>( + Struct {Field::create_field<TYPE_MAP>(map), Field::create_field<TYPE_STRING>("17")})); + column->insert_default(); + Block block; + block.insert({std::move(column), type, "nested"}); + std::shared_ptr<arrow::Schema> schema; + ASSERT_TRUE(ArrowFlightSchemaConvertor("UTC").get_arrow_schema_from_block(block, &schema).ok()); + ArrowFlightArrowBlockConvertor converter(schema, cctz::utc_time_zone()); + std::shared_ptr<arrow::RecordBatch> batch; + auto status = converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch); + ASSERT_TRUE(status.ok()) << status; + ASSERT_TRUE(batch->ValidateFull().ok()); + std::string serialized; + ASSERT_TRUE(serialize_record_batch(*batch, &serialized).ok()); + auto source = std::make_shared<arrow::io::BufferReader>(arrow::Buffer::FromString(serialized)); + auto reader = arrow::ipc::RecordBatchStreamReader::Open(source).ValueOrDie(); + ASSERT_TRUE(reader->ReadNext(&batch).ok()); + const auto& structure = static_cast<const arrow::StructArray&>(*batch->column(0)); + EXPECT_TRUE(structure.IsNull(1)); + const auto& maps = static_cast<const arrow::MapArray&>(*structure.field(0)); + const auto& arrow_map_type = static_cast<const arrow::MapType&>(*maps.type()); + expect_logical_type(arrow_map_type.key_field(), "LARGEINT"); + expect_logical_type(arrow_map_type.item_field(), "LARGEINT"); + EXPECT_FALSE(arrow_map_type.key_field()->nullable()); + EXPECT_EQ("170141183460469231731687303715884105727", + static_cast<const arrow::StringArray&>(*maps.keys()).GetString(0)); + EXPECT_EQ("-170141183460469231731687303715884105728", + static_cast<const arrow::StringArray&>(*maps.items()).GetString(0)); + EXPECT_EQ("17", static_cast<const arrow::StringArray&>(*structure.field(1)).GetString(0)); + EXPECT_EQ(nullptr, structure.type()->field(1)->metadata()); + expect_legacy_batch(block, batch); +} + +} // namespace +} // namespace doris diff --git a/be/test/format/table/iceberg/arrow_schema_util_test.cpp b/be/test/format/table/iceberg/arrow_schema_util_test.cpp index d8241812639..063302a3252 100644 --- a/be/test/format/table/iceberg/arrow_schema_util_test.cpp +++ b/be/test/format/table/iceberg/arrow_schema_util_test.cpp @@ -56,7 +56,7 @@ TEST(ArrowSchemaUtilTest, test_simple_field) { std::vector<std::shared_ptr<arrow::Field>> fields; Status st; - st = ArrowSchemaUtil::convert(&schema, "utc", fields); + st = IcebergArrowSchemaConvertor(schema, "utc").convert_fields(fields); EXPECT_TRUE(st.ok()); EXPECT_EQ(2, fields.size()); EXPECT_EQ("field1", fields[0]->name()); @@ -114,7 +114,7 @@ TEST(ArrowSchemaUtilTest, test_stuct_field) { std::vector<std::shared_ptr<arrow::Field>> fields; Status st; - st = ArrowSchemaUtil::convert(schema.get(), "utc", fields); + st = IcebergArrowSchemaConvertor(*schema, "utc").convert_fields(fields); EXPECT_TRUE(st.ok()); EXPECT_EQ(1, fields.size()); EXPECT_EQ("st_col", fields[0]->name()); @@ -166,7 +166,7 @@ TEST(ArrowSchemaUtilTest, test_map_field) { std::vector<std::shared_ptr<arrow::Field>> fields; Status st; - st = ArrowSchemaUtil::convert(schema.get(), "utc", fields); + st = IcebergArrowSchemaConvertor(*schema, "utc").convert_fields(fields); EXPECT_TRUE(st.ok()); EXPECT_EQ(1, fields.size()); EXPECT_EQ("map_col", fields[0]->name()); @@ -213,7 +213,7 @@ TEST(ArrowSchemaUtilTest, test_list_field) { std::vector<std::shared_ptr<arrow::Field>> fields; Status st; - st = ArrowSchemaUtil::convert(schema.get(), "utc", fields); + st = IcebergArrowSchemaConvertor(*schema, "utc").convert_fields(fields); EXPECT_TRUE(st.ok()); EXPECT_EQ(1, fields.size()); EXPECT_EQ("list_col", fields[0]->name()); @@ -237,7 +237,7 @@ TEST(ArrowSchemaUtilTest, test_binary_field_types) { Schema schema(1, std::move(nested_fields)); std::vector<std::shared_ptr<arrow::Field>> fields; - Status st = ArrowSchemaUtil::convert(&schema, "utc", fields); + Status st = IcebergArrowSchemaConvertor(schema, "utc").convert_fields(fields); ASSERT_TRUE(st.ok()) << st; ASSERT_EQ(4, fields.size()); @@ -260,7 +260,7 @@ TEST(ArrowSchemaUtilTest, test_variant_field) { Schema schema(1, std::move(nested_fields)); std::vector<std::shared_ptr<arrow::Field>> fields; - Status st = ArrowSchemaUtil::convert(&schema, "utc", fields); + Status st = IcebergArrowSchemaConvertor(schema, "utc").convert_fields(fields); ASSERT_TRUE(st.ok()) << st; ASSERT_EQ(1, fields.size()); EXPECT_EQ("21", fields[0]->metadata()->Get(pfid).ValueUnsafe()); @@ -324,7 +324,7 @@ TEST(ArrowSchemaUtilTest, test_nested_variant_fields) { std::unique_ptr<Schema> schema = SchemaParser::from_json(schema_json); std::vector<std::shared_ptr<arrow::Field>> fields; - Status status = ArrowSchemaUtil::convert(schema.get(), "UTC", fields); + Status status = IcebergArrowSchemaConvertor(*schema, "UTC").convert_fields(fields); ASSERT_TRUE(status.ok()) << status; ASSERT_EQ(3, fields.size()); @@ -387,7 +387,7 @@ TEST(ArrowSchemaUtilTest, test_parquet_filed_id) { Schema schema(1, std::move(nested_fields)); std::vector<std::shared_ptr<arrow::Field>> fields; - st = ArrowSchemaUtil::convert(&schema, "utc", fields); + st = IcebergArrowSchemaConvertor(schema, "utc").convert_fields(fields); auto arrow_schema = arrow::schema(fields); // create arrow table diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/ResultSink.java b/fe/fe-core/src/main/java/org/apache/doris/planner/ResultSink.java index 25e72ed7598..b7d7dd4d876 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/ResultSink.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/ResultSink.java @@ -73,6 +73,9 @@ public class ResultSink extends DataSink { tResultSink.setFetchOption(fetchOption); } tResultSink.setType(resultSinkType); + if (resultSinkType == TResultSinkType.ARROW_FLIGHT_PROTOCOL) { + tResultSink.setEnableArrowTypeMetadata(true); + } result.setResultSink(tResultSink); return result; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java index 2296986c535..5942424750c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java @@ -40,6 +40,7 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ArrowStreamReader; +import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -158,7 +159,13 @@ public class FlightSqlConnectProcessor extends ConnectProcessor implements AutoC String.format("Schema size %s' is not equal to arrow field size %s, queryId: %s.", fieldVectors.size(), resultOutputExprs.size(), DebugUtil.printId(tid))); } - schema = root.getSchema(); + List<Field> fields = new ArrayList<>(); + for (int i = 0; i < resultOutputExprs.size(); i++) { + fields.add(FlightSqlSchemaHelper.withDorisTypeMetadata( + root.getSchema().getFields().get(i), resultOutputExprs.get(i).getType())); + } + // Compare and advertise the same complete metadata regardless of BE upgrade order. + schema = new Schema(fields, root.getSchema().getCustomMetadata()); if (arrowSchema == null) { arrowSchema = schema; } else if (!arrowSchema.equals(schema)) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlSchemaHelper.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlSchemaHelper.java index 9d91153ebd6..6af37014132 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlSchemaHelper.java +++ b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlSchemaHelper.java @@ -17,8 +17,12 @@ package org.apache.doris.service.arrowflight; +import org.apache.doris.catalog.ArrayType; import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.MapType; import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.StructType; +import org.apache.doris.catalog.Type; import org.apache.doris.datasource.CatalogIf; import org.apache.doris.qe.ConnectContext; import org.apache.doris.service.ExecuteEnv; @@ -165,6 +169,47 @@ public class FlightSqlSchemaHelper { } } + static Field withDorisTypeMetadata(Field field, Type type) { + List<Field> children = new ArrayList<>(field.getChildren()); + if (type.isArrayType()) { + children.set(0, withDorisTypeMetadata(children.get(0), ((ArrayType) type).getItemType())); + } else if (type.isMapType()) { + Field entries = children.get(0); + List<Field> pair = new ArrayList<>(entries.getChildren()); + pair.set(0, withDorisTypeMetadata(pair.get(0), ((MapType) type).getKeyType())); + pair.set(1, withDorisTypeMetadata(pair.get(1), ((MapType) type).getValueType())); + children.set(0, new Field(entries.getName(), entries.getFieldType(), pair)); + } else if (type.isStructType()) { + StructType struct = (StructType) type; + for (int i = 0; i < children.size(); i++) { + children.set(i, withDorisTypeMetadata(children.get(i), struct.getFields().get(i).getType())); + } + } + String marker = null; + switch (type.getPrimitiveType()) { + case LARGEINT: + case IPV4: + case IPV6: + case VARIANT: + marker = type.getPrimitiveType().name(); + break; + case JSONB: + marker = "JSON"; + break; + default: + break; + } + FieldType fieldType = field.getFieldType(); + if (marker != null && !field.getMetadata().containsKey("doris_type")) { + Map<String, String> metadata = new HashMap<>(field.getMetadata()); + metadata.put("doris_type", marker); + fieldType = new FieldType(field.isNullable(), field.getType(), field.getDictionary(), metadata); + } + // Old BEs omit these markers. Fill only missing ones from the planned Doris type; + // preserve conflicting markers and all physical properties for the strict schema comparison. + return new Field(field.getName(), fieldType, children); + } + private static ArrowType columnDescToArrowType(final TColumnDesc desc) { PrimitiveType primitiveType = PrimitiveType.fromThrift(desc.getColumnType()); Integer precision = desc.isSetColumnPrecision() ? desc.getColumnPrecision() : null; diff --git a/fe/fe-core/src/test/java/org/apache/doris/planner/ResultSinkMetadataTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/ResultSinkMetadataTest.java new file mode 100644 index 00000000000..d41c1062155 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/planner/ResultSinkMetadataTest.java @@ -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. + +package org.apache.doris.planner; + +import org.apache.doris.thrift.TResultSink; +import org.apache.doris.thrift.TResultSinkType; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class ResultSinkMetadataTest { + @Test + void onlyUpgradedFlightPlannerOptsIntoExtendedMetadata() { + Assertions.assertFalse(new TResultSink().isEnableArrowTypeMetadata()); + TResultSink flight = new ResultSink(new PlanNodeId(0), TResultSinkType.ARROW_FLIGHT_PROTOCOL) + .toThrift().getResultSink(); + Assertions.assertTrue(flight.isSetEnableArrowTypeMetadata()); + Assertions.assertTrue(flight.isEnableArrowTypeMetadata()); + TResultSink mysql = new ResultSink(new PlanNodeId(0), TResultSinkType.MYSQL_PROTOCOL) + .toThrift().getResultSink(); + Assertions.assertFalse(mysql.isSetEnableArrowTypeMetadata()); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessorSchemaTest.java b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessorSchemaTest.java new file mode 100644 index 00000000000..b7f6725478b --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessorSchemaTest.java @@ -0,0 +1,167 @@ +// 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. + +package org.apache.doris.service.arrowflight; + +import org.apache.doris.analysis.Expr; +import org.apache.doris.catalog.ArrayType; +import org.apache.doris.catalog.MapType; +import org.apache.doris.catalog.StructField; +import org.apache.doris.catalog.StructType; +import org.apache.doris.catalog.Type; +import org.apache.doris.proto.InternalService.PFetchArrowFlightSchemaResult; +import org.apache.doris.proto.Types.PStatus; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.rpc.BackendServiceProxy; +import org.apache.doris.service.arrowflight.results.FlightSqlEndpointsLocation; +import org.apache.doris.thrift.TNetworkAddress; +import org.apache.doris.thrift.TUniqueId; + +import com.google.protobuf.ByteString; +import org.apache.arrow.vector.ipc.WriteChannel; +import org.apache.arrow.vector.ipc.message.MessageSerializer; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.io.ByteArrayOutputStream; +import java.nio.channels.Channels; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; + +class FlightSqlConnectProcessorSchemaTest { + private static Field field(String name, ArrowType type, boolean nullable, String marker, Field... children) { + Map<String, String> metadata = marker == null ? Collections.emptyMap() + : Collections.singletonMap("doris_type", marker); + return new Field(name, new FieldType(nullable, type, null, metadata), Arrays.asList(children)); + } + + private static Schema schema(Field... fields) { + return new Schema(Arrays.asList(fields), Collections.singletonMap("schema_key", "schema_value")); + } + + private static Schema nestedSchema(boolean annotated) { + return schema( + field("a", new ArrowType.List(), true, null, + field("item", new ArrowType.Utf8(), true, annotated ? "LARGEINT" : null)), + field("m", new ArrowType.Map(false), true, null, + field("entries", new ArrowType.Struct(), false, null, + field("key", new ArrowType.Utf8(), false, annotated ? "LARGEINT" : null), + field("value", new ArrowType.Struct(), true, null, + field("ip4", new ArrowType.Int(32, true), true, annotated ? "IPV4" : null), + field("ip6", new ArrowType.Utf8(), true, annotated ? "IPV6" : null), + field("json", new ArrowType.Utf8(), true, annotated ? "JSON" : null), + field("variant", new ArrowType.Utf8(), true, annotated ? "VARIANT" : null), + field("text", new ArrowType.Utf8(), true, null)))), + field("json", new ArrowType.Utf8(), true, annotated ? "JSON" : null), + field("variant", new ArrowType.Utf8(), true, annotated ? "VARIANT" : null)); + } + + private static List<Type> resultTypes() { + return Arrays.asList(new ArrayType(Type.LARGEINT), + new MapType(Type.LARGEINT, new StructType( + new StructField("ip4", Type.IPV4), new StructField("ip6", Type.IPV6), + new StructField("json", Type.JSONB), new StructField("variant", Type.VARIANT), + new StructField("text", Type.STRING))), Type.JSONB, Type.VARIANT); + } + + private static Schema fetch(List<Type> types, Schema... schemas) throws Exception { + List<FlightSqlEndpointsLocation> endpoints = new ArrayList<>(); + List<CompletableFuture<PFetchArrowFlightSchemaResult>> responses = new ArrayList<>(); + for (int i = 0; i < schemas.length; i++) { + ArrayList<Expr> exprs = new ArrayList<>(); + for (Type type : types) { + Expr expr = Mockito.mock(Expr.class); + Mockito.when(expr.getType()).thenReturn(type); + exprs.add(expr); + } + endpoints.add(new FlightSqlEndpointsLocation(new TUniqueId(1, i), + new TNetworkAddress("localhost", 10000 + i), + new TNetworkAddress("localhost", 11000 + i), exprs)); + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(bytes)), schemas[i]); + responses.add(CompletableFuture.completedFuture(PFetchArrowFlightSchemaResult.newBuilder() + .setStatus(PStatus.newBuilder().setStatusCode(0)) + .setSchema(ByteString.copyFrom(bytes.toByteArray())).build())); + } + ConnectContext context = Mockito.mock(ConnectContext.class); + Mockito.when(context.getFlightSqlEndpointsLocations()).thenReturn(endpoints); + BackendServiceProxy proxy = Mockito.mock(BackendServiceProxy.class); + AtomicInteger next = new AtomicInteger(); + Mockito.when(proxy.fetchArrowFlightSchema(Mockito.any(), Mockito.any())) + .thenAnswer(invocation -> responses.get(next.getAndIncrement())); + try (MockedStatic<BackendServiceProxy> singleton = Mockito.mockStatic(BackendServiceProxy.class); + FlightSqlConnectProcessor processor = new FlightSqlConnectProcessor(context)) { + singleton.when(BackendServiceProxy::getInstance).thenReturn(proxy); + processor.fetchArrowFlightSchema(1000); + return processor.getArrowSchema(); + } + } + + @Test + void mixedVersionsAdvertiseCompleteMetadataInEitherOrder() throws Exception { + Schema oldSchema = nestedSchema(false); + Schema newSchema = nestedSchema(true); + Assertions.assertEquals(newSchema, fetch(resultTypes(), oldSchema, newSchema)); + Assertions.assertEquals(newSchema, fetch(resultTypes(), newSchema, oldSchema)); + } + + @Test + void oldOnlyAndNewOnlyAdvertiseTheSameSchema() throws Exception { + Assertions.assertEquals(nestedSchema(true), fetch(resultTypes(), nestedSchema(false))); + Assertions.assertEquals(nestedSchema(true), fetch(resultTypes(), nestedSchema(true), nestedSchema(true))); + } + + private static void assertSchemaMismatch(Schema first, Schema second) { + RuntimeException failure = Assertions.assertThrows(RuntimeException.class, + () -> fetch(Collections.singletonList(Type.JSONB), first, second)); + Assertions.assertTrue(failure.getCause().getMessage() + .startsWith("The schema returned by results BE is different")); + } + + @Test + void stillRejectsPhysicalDifferences() { + Schema expected = schema(field("value", new ArrowType.Utf8(), true, "JSON")); + for (Field incompatible : Arrays.asList( + field("value", new ArrowType.Int(32, true), true, "JSON"), + field("renamed", new ArrowType.Utf8(), true, "JSON"), + field("value", new ArrowType.Utf8(), false, "JSON"))) { + assertSchemaMismatch(expected, schema(incompatible)); + } + } + + @Test + void stillRejectsConflictingOrUnrelatedMetadata() { + Schema expected = schema(field("value", new ArrowType.Utf8(), true, "JSON")); + Schema wrongMarker = schema(field("value", new ArrowType.Utf8(), true, "VARIANT")); + Field extraMetadata = new Field("value", new FieldType(true, new ArrowType.Utf8(), null, + Collections.singletonMap("other_key", "other_value")), Collections.emptyList()); + assertSchemaMismatch(expected, wrongMarker); + assertSchemaMismatch(expected, schema(extraMetadata)); + assertSchemaMismatch(expected, new Schema(expected.getFields())); + } +} diff --git a/gensrc/thrift/DataSinks.thrift b/gensrc/thrift/DataSinks.thrift index c873e90529c..321bdf94469 100644 --- a/gensrc/thrift/DataSinks.thrift +++ b/gensrc/thrift/DataSinks.thrift @@ -232,6 +232,8 @@ struct TResultSink { 1: optional TResultSinkType type; 2: optional TResultFileSinkOptions file_options; // deprecated 3: optional TFetchOption fetch_option; + // Absent for old FEs, whose Flight schema comparison requires legacy field metadata. + 4: optional bool enable_arrow_type_metadata = false; } struct TResultFileSink { diff --git a/samples/arrow-flight-sql/python/README.md b/samples/arrow-flight-sql/python/README.md index e729054b813..7a7586940a1 100644 --- a/samples/arrow-flight-sql/python/README.md +++ b/samples/arrow-flight-sql/python/README.md @@ -35,6 +35,28 @@ under the License. Section 6.1 of https://github.com/apache/doris/issues/25514 is the performance test results of the doris arrow flight sql using python. +## Logical type metadata + +Some Doris types share an Arrow storage type with ordinary strings or integers. +The `doris_type` field metadata identifies `LARGEINT`, `IPV4`, `IPV6`, `JSON`, and +`VARIANT`, including fields nested in arrays, maps, and structs. Consumers should +inspect each nested Arrow field instead of inferring a type from its value. + +LARGEINT retains its Arrow string encoding, including the full signed 128-bit +range. PyArrow does not automatically convert custom metadata into Python types; +a client can use `doris_type=LARGEINT` to safely convert that field's non-NULL +values with `int(value)`, while leaving ordinary STRING fields unchanged. + +During rolling upgrades, queries planned by older FEs retain the legacy metadata +layout. Upgraded FEs advertise complete logical type metadata in the FlightInfo +schema even when some BEs are older. An older BE's DoGet stream may still omit +these markers; use the FlightInfo schema as the logical type reference until the +upgrade finishes. No session setting is required. + +Run `python test_nested_type_metadata.py` with `DORIS_FLIGHT_SQL_URI` and optional +`DORIS_USER` / `DORIS_PASSWORD` to verify the metadata and values against a cluster. +The tests execute only read-only queries. + # Notes For more details, refer to [Python Usage] in the document https://doris.apache.org/zh-CN/docs/dev/db-connect/arrow-flight-sql-connect diff --git a/samples/arrow-flight-sql/python/test_nested_type_metadata.py b/samples/arrow-flight-sql/python/test_nested_type_metadata.py new file mode 100644 index 00000000000..f918cf5702c --- /dev/null +++ b/samples/arrow-flight-sql/python/test_nested_type_metadata.py @@ -0,0 +1,108 @@ +# 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. + +"""Run with DORIS_FLIGHT_SQL_URI and optional DORIS_USER/DORIS_PASSWORD. + +LARGEINT retains its string storage encoding. PyArrow does not interpret custom +field metadata automatically; applications can use doris_type to distinguish +these strings from text and convert them to Python integers without losing range. +""" + +import os +import unittest + +import adbc_driver_flightsql.dbapi as flight_sql +import adbc_driver_manager +import pyarrow as pa + + [email protected](os.getenv("DORIS_FLIGHT_SQL_URI"), "Set DORIS_FLIGHT_SQL_URI to run against Doris") +class NestedTypeMetadataTest(unittest.TestCase): + def setUp(self): + self.connection = flight_sql.connect( + uri=os.environ["DORIS_FLIGHT_SQL_URI"], + db_kwargs={ + adbc_driver_manager.DatabaseOptions.USERNAME.value: os.getenv("DORIS_USER", "root"), + adbc_driver_manager.DatabaseOptions.PASSWORD.value: os.getenv("DORIS_PASSWORD", ""), + }, + ) + self.addCleanup(self.connection.close) + self.cursor = self.connection.cursor() + self.addCleanup(self.cursor.close) + + def assert_logical_type(self, field, name): + self.assertEqual((field.metadata or {}).get(b"doris_type"), name.encode()) + + def query(self, sql): + self.cursor.execute(sql) + table = self.cursor.fetch_arrow_table() + table.validate(full=True) + return table + + def test_largeint_metadata_and_full_range(self): + table = self.query(""" + SELECT CAST('170141183460469231731687303715884105727' AS LARGEINT) AS scalar_value, + named_struct('number', CAST(17 AS LARGEINT), 'text', '17') AS struct_value, + array(CAST('-170141183460469231731687303715884105728' AS LARGEINT), + CAST(NULL AS LARGEINT)) AS array_value, + map(CAST(17 AS LARGEINT), CAST(19 AS LARGEINT)) AS map_value + """) + self.assert_logical_type(table.schema.field("scalar_value"), "LARGEINT") + structure = table.schema.field("struct_value").type + self.assert_logical_type(structure.field("number"), "LARGEINT") + self.assertNotIn(b"doris_type", structure.field("text").metadata or {}) + self.assert_logical_type(table.schema.field("array_value").type.value_field, "LARGEINT") + mapping = table.schema.field("map_value").type + self.assert_logical_type(mapping.key_field, "LARGEINT") + self.assert_logical_type(mapping.item_field, "LARGEINT") + self.assertFalse(mapping.key_field.nullable) + self.assertTrue(pa.types.is_string(table.schema.field("scalar_value").type)) + row = table.to_pylist()[0] + self.assertEqual(row["scalar_value"], str(2**127 - 1)) + self.assertEqual(row["array_value"], [str(-(2**127)), None]) + self.assertEqual(row["struct_value"], {"number": "17", "text": "17"}) + self.assertEqual(row["map_value"], [("17", "19")]) + + def test_ipv4_metadata(self): + table = self.query(""" + SELECT array(CAST('192.0.2.1' AS IPV4)) AS ip4 + """) + self.assert_logical_type(table.schema.field("ip4").type.value_field, "IPV4") + + def test_ipv6_metadata(self): + table = self.query(""" + SELECT named_struct('address', CAST('2001:db8::1' AS IPV6)) AS ip6 + """) + self.assert_logical_type(table.schema.field("ip6").type.field("address"), "IPV6") + + def test_json_metadata(self): + # ARRAY/MAP constructors reject JSON; ARRAY_REPEAT preserves its element type. + table = self.query(""" + SELECT array_repeat(CAST('{"n":1}' AS JSON), 1) AS json_value + """) + self.assert_logical_type(table.schema.field("json_value").type.value_field, "JSON") + + def test_variant_metadata(self): + # ARRAY rejects legacy VARIANT; ARRAY_REPEAT also supports the default mode. + table = self.query(""" + SELECT array_repeat(CAST('{"n":1}' AS VARIANT), 1) AS variant_value + """) + self.assert_logical_type(table.schema.field("variant_value").type.value_field, "VARIANT") + + +if __name__ == "__main__": + unittest.main() --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
