SteNicholas commented on code in PR #376: URL: https://github.com/apache/paimon-cpp/pull/376#discussion_r4255423490
########## src/paimon/format/vortex/vortex_file_batch_reader.cpp: ########## @@ -0,0 +1,457 @@ +/* + * 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 "paimon/format/vortex/vortex_file_batch_reader.h" + +#include <algorithm> +#include <limits> +#include <string> +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/predicate/predicate_filter.h" +#include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/math.h" +#include "paimon/core/utils/nested_projection_utils.h" +#include "paimon/format/vortex/vortex_ffi_util.h" +#include "paimon/fs/file_system.h" + +namespace paimon::vortex { + +namespace { + +// Vortex's Arrow export represents strings and binaries with the view types (StringView / +// BinaryView). paimon schemas use the standard string/binary types, so map the view types back, +// recursing through nested types. Types that contain no views are returned unchanged. +std::shared_ptr<arrow::DataType> NormalizeViewType(const std::shared_ptr<arrow::DataType>& type) { + switch (type->id()) { + case arrow::Type::STRING_VIEW: + return arrow::utf8(); + case arrow::Type::BINARY_VIEW: + return arrow::binary(); + case arrow::Type::STRUCT: { + arrow::FieldVector fields; + fields.reserve(type->num_fields()); + for (const std::shared_ptr<arrow::Field>& field : type->fields()) { + fields.push_back(field->WithType(NormalizeViewType(field->type()))); + } + return arrow::struct_(fields); + } + case arrow::Type::LIST: + return arrow::list(type->field(0)->WithType(NormalizeViewType(type->field(0)->type()))); + case arrow::Type::FIXED_SIZE_LIST: + return arrow::fixed_size_list( + type->field(0)->WithType(NormalizeViewType(type->field(0)->type())), + checked_cast<const arrow::FixedSizeListType&>(*type).list_size()); + default: + return type; + } +} + +std::shared_ptr<arrow::Schema> NormalizeViewSchema(const std::shared_ptr<arrow::Schema>& schema) { + arrow::FieldVector fields; + fields.reserve(schema->num_fields()); + for (const std::shared_ptr<arrow::Field>& field : schema->fields()) { + fields.push_back(field->WithType(NormalizeViewType(field->type()))); + } + return arrow::schema(fields, schema->metadata()); +} + +// Rebuilds `array` with every StringView/BinaryView column (which Vortex's Arrow export produces) +// converted to the standard string/binary type. Arrow 17 has no cast kernel between the view and +// standard types, so the leaves are rebuilt with builders and the parents reconstructed, recursing +// through structs and lists. Arrays with no view-typed data are returned unchanged. `array` must be +// offset-normalized (offset 0) so the rebuilt parents and children stay consistent. +Result<std::shared_ptr<arrow::Array>> NormalizeViewArray(const std::shared_ptr<arrow::Array>& array, + arrow::MemoryPool* pool) { + // TODO(zhouyc-ali): avoid the per-row copy by exporting standard utf8/binary + switch (array->type_id()) { + case arrow::Type::STRING_VIEW: { + const auto& view = checked_cast<const arrow::StringViewArray&>(*array); + arrow::StringBuilder builder(pool); + for (int64_t i = 0; i < view.length(); ++i) { + if (view.IsNull(i)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(view.GetView(i))); + } + } + std::shared_ptr<arrow::Array> out; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&out)); + return out; + } + case arrow::Type::BINARY_VIEW: { + const auto& view = checked_cast<const arrow::BinaryViewArray&>(*array); + arrow::BinaryBuilder builder(pool); + for (int64_t i = 0; i < view.length(); ++i) { + if (view.IsNull(i)) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); + } else { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(view.GetView(i))); + } + } + std::shared_ptr<arrow::Array> out; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&out)); + return out; + } + case arrow::Type::STRUCT: { + const auto& struct_array = checked_cast<const arrow::StructArray&>(*array); + arrow::ArrayVector children; + arrow::FieldVector fields; + children.reserve(struct_array.num_fields()); + fields.reserve(struct_array.num_fields()); + bool changed = false; + for (int32_t i = 0; i < struct_array.num_fields(); ++i) { + const std::shared_ptr<arrow::Array>& child = struct_array.field(i); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> normalized_child, + NormalizeViewArray(child, pool)); + changed = changed || normalized_child.get() != child.get(); + fields.push_back(array->type()->field(i)->WithType(normalized_child->type())); + children.push_back(std::move(normalized_child)); + } + if (!changed) { + return array; + } + return std::make_shared<arrow::StructArray>( + arrow::struct_(fields), struct_array.length(), children, struct_array.null_bitmap(), + struct_array.null_count(), struct_array.offset()); + } + case arrow::Type::LIST: { + const auto& list = checked_cast<const arrow::ListArray&>(*array); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> values, + NormalizeViewArray(list.values(), pool)); + if (values.get() == list.values().get()) { + return array; + } + return std::make_shared<arrow::ListArray>( + arrow::list(array->type()->field(0)->WithType(values->type())), list.length(), + list.value_offsets(), values, list.null_bitmap(), list.null_count(), list.offset()); + } + case arrow::Type::FIXED_SIZE_LIST: { + const auto& list = checked_cast<const arrow::FixedSizeListArray&>(*array); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> values, + NormalizeViewArray(list.values(), pool)); + if (values.get() == list.values().get()) { + return array; + } + const auto& fsl_type = checked_cast<const arrow::FixedSizeListType&>(*array->type()); + return std::make_shared<arrow::FixedSizeListArray>( + arrow::fixed_size_list(array->type()->field(0)->WithType(values->type()), + fsl_type.list_size()), + list.length(), values, list.null_bitmap(), list.null_count(), list.offset()); + } + default: + return array; + } +} + +} // namespace + +VortexFileBatchReader::VortexFileBatchReader( + const std::shared_ptr<InputStream>& input, int32_t batch_size, + std::shared_ptr<VortexInputContext> input_context, VxSessionPtr session, + VxDataSourcePtr data_source, VxScanPtr scan, const std::shared_ptr<arrow::Schema>& file_schema, + const std::shared_ptr<arrow::DataType>& struct_type, uint64_t total_rows, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) + : input_(input), + batch_size_(batch_size), + input_context_(std::move(input_context)), + session_(std::move(session)), + data_source_(std::move(data_source)), + scan_(std::move(scan)), + file_schema_(NormalizeViewSchema(file_schema)), + struct_type_(struct_type), + total_rows_(total_rows), + arrow_pool_(arrow_pool), + metrics_(std::make_shared<MetricsImpl>()) {} + +Result<std::unique_ptr<VortexFileBatchReader>> VortexFileBatchReader::Create( + const std::shared_ptr<InputStream>& input, int32_t batch_size, + const std::shared_ptr<arrow::MemoryPool>& arrow_pool) { + if (input == nullptr || arrow_pool == nullptr || batch_size <= 0) { + return Status::Invalid( + "Vortex reader requires non-null input and Arrow memory pool, and positive batch size"); + } + PAIMON_ASSIGN_OR_RAISE(int64_t signed_length, input->Length()); + PAIMON_RETURN_NOT_OK(ValidateValueNonNegative(signed_length, "Vortex input length")); + auto length = static_cast<uint64_t>(signed_length); + auto input_context = std::make_shared<VortexInputContext>(input); + + VxSessionPtr session(vx_session_new(), vx_session_free); + if (session == nullptr) { + return Status::IOError("failed to create Vortex session"); + } + vx_error* error = nullptr; + VxDataSourcePtr data_source( + vx_data_source_new_callback(session.get(), VortexInputContext::MakeCallbacks(input_context), + length, &error), + vx_data_source_free); + if (data_source == nullptr) { + return VortexCallbackError("open Vortex callback data source", error, + input_context->GetCallbackStatus()); + } + + // The data source dtype is available without consuming a scan. + // NOTE(vortex 0.75): vx_data_source_dtype returns a BORROWED pointer (no refcount bump), so it + // must NOT be freed here: freeing decrements the DType Arc and segfaults later in the scan. + // Upstream >=0.77 returns an owned clone that MUST be freed with vx_dtype_free. + const vx_dtype* dtype = vx_data_source_dtype(data_source.get()); + if (dtype == nullptr) { + return Status::IOError("failed to read Vortex data source dtype"); + } + ::ArrowSchema ffi_schema = {}; + error = nullptr; + if (vx_dtype_to_arrow_schema(dtype, &ffi_schema, &error) != 0) { + return VortexFfiError("convert Vortex dtype to Arrow schema", error); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema> file_schema, + arrow::ImportSchema(&ffi_schema)); + // Built from the raw view-typed schema on purpose: NextBatch imports each scanned batch as this + // exact type, while the constructor normalizes its own copy kept as file_schema_. + std::shared_ptr<arrow::DataType> struct_type = arrow::struct_(file_schema->fields()); + + vx_estimate row_count{}; + vx_data_source_get_row_count(data_source.get(), &row_count); + if (row_count.type == VX_ESTIMATE_UNKNOWN) { + return Status::Invalid("Vortex data source did not report a row count"); + } + const uint64_t total_rows = row_count.estimate; + + // The scan is created lazily on first read (and re-created by SetReadSchema), since a Vortex + // scan may be consumed only once. + return std::unique_ptr<VortexFileBatchReader>(new VortexFileBatchReader( + input, batch_size, std::move(input_context), std::move(session), std::move(data_source), + VxScanPtr(nullptr, vx_scan_free), file_schema, struct_type, total_rows, arrow_pool)); +} + +VortexFileBatchReader::~VortexFileBatchReader() { + CloseInternal(); +} + +Result<bool> VortexFileBatchReader::OpenNextPartitionStream() { + if (scan_ == nullptr) { + vx_error* error = nullptr; + // Scan in storage order: Vortex defaults to ordered=false and may emit chunks out of order, + // which would misalign the physical row positions NextBatch assigns. + vx_scan_options options{}; // zero-init: all columns, no filter/row-range/selection/limit + options.ordered = true; + VxScanPtr scan( + vx_data_source_scan(data_source_.get(), &options, /*estimate=*/nullptr, &error), + vx_scan_free); + if (scan == nullptr) { + return VortexCallbackError("create Vortex scan", error, + input_context_->GetCallbackStatus()); + } + scan_ = std::move(scan); + } + vx_error* error = nullptr; + VxPartitionPtr partition(vx_scan_next_partition(scan_.get(), &error), vx_partition_free); + if (error != nullptr) { + return VortexCallbackError("advance Vortex scan partition", error, + input_context_->GetCallbackStatus()); + } + if (partition == nullptr) { + return false; // Scan exhausted. + } + // vx_partition_scan_arrow consumes the partition (even on error), so release ownership here. + ::ArrowArrayStream stream{}; + error = nullptr; + if (vx_partition_scan_arrow(session_.get(), partition.release(), &stream, &error) != 0) { + return VortexCallbackError("scan Vortex partition to Arrow", error, + input_context_->GetCallbackStatus()); + } + current_stream_ = stream; + stream_active_ = true; + return true; +} + +Result<std::shared_ptr<arrow::Array>> VortexFileBatchReader::ReadNextArray() { + while (true) { + if (!stream_active_) { + PAIMON_ASSIGN_OR_RAISE(bool opened, OpenNextPartitionStream()); + if (!opened) { + return std::shared_ptr<arrow::Array>(); // End of scan. + } + } + ::ArrowArray ffi_array = {}; + const int32_t rc = current_stream_.get_next(¤t_stream_, &ffi_array); + if (rc != 0) { + // get_last_error returns a pointer owned by the stream; the Arrow C stream contract + // invalidates it once the stream is released, so copy it before ReleaseStream(). + const char* message = current_stream_.get_last_error(¤t_stream_); + std::string message_copy = message == nullptr ? "unknown" : message; + // A read failure surfaces here as an opaque stream error, so prefer the status the IO + // callback stashed. + Status callback_status = input_context_->GetCallbackStatus(); + ReleaseStream(); + if (!callback_status.ok()) { + return callback_status.WithMessage("read Vortex batch", ": ", + callback_status.message()); + } + return Status::IOError("Vortex Arrow stream error: ", message_copy); + } + if (ffi_array.release == nullptr) { + ReleaseStream(); // Partition exhausted; try the next one. + continue; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array, + arrow::ImportArray(&ffi_array, struct_type_)); + if (array->length() == 0) { + continue; // Skip empty batches. + } + return array; + } +} + +Result<std::shared_ptr<arrow::Array>> VortexFileBatchReader::AlignToReadSchema( + const std::shared_ptr<arrow::Array>& array) const { + if (read_schema_ == nullptr) { + return array; + } + // TODO(zhouyc-ali): push top-level column pruning into the Vortex scan + return NestedProjectionUtils::AlignArrayToReadType( + array, arrow::struct_(read_schema_->fields()), arrow_pool_.get()); Review Comment: **Normalize timestamp timezones before schema alignment.** Vortex preserves the stored timestamp timezone, while Paimon reconstructs `TIMESTAMP_LTZ` using the reader's local timezone. `AlignArrayToReadType` rejects that difference: UTC data requested as `Asia/Shanghai` fails with `unsupported leaf type change`. This prevents reading the Java fixture's `TIMESTAMP_LTZ` column outside UTC; its compatibility test masks the issue with `TimezoneGuard("UTC")`. Please retype timezone-bearing timestamps without changing their epoch values before alignment, including nested columns, and test a non-UTC reader. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
