Gabriel39 commented on code in PR #63444: URL: https://github.com/apache/doris/pull/63444#discussion_r3286866069
########## be/src/format/new_parquet/column_reader.cpp: ########## @@ -0,0 +1,909 @@ +// 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/new_parquet/column_reader.h" + +#include <arrow/array/array_binary.h> +#include <parquet/api/reader.h> +#include <parquet/api/schema.h> + +#include <algorithm> +#include <cstddef> +#include <cstdint> +#include <exception> +#include <memory> +#include <string> +#include <type_traits> +#include <utility> +#include <vector> + +#include "core/column/column.h" +#include "core/column/column_decimal.h" +#include "core/column/column_struct.h" +#include "core/data_type/data_type_nullable.h" +#include "core/data_type/data_type_number.h" +#include "core/value/vdatetime_value.h" +#include "format/new_parquet/parquet_column_schema.h" + +namespace doris::parquet { +namespace { + +class PrimitiveColumnReader final : public ParquetColumnReader { +public: + PrimitiveColumnReader(int file_column_id, const ::parquet::ColumnDescriptor* descriptor, + ParquetTypeDescriptor type_descriptor, DataTypePtr type, std::string name, + std::shared_ptr<::parquet::internal::RecordReader> record_reader) + : _file_column_id(file_column_id), + _parquet_column_ordinal(file_column_id), + _descriptor(descriptor), + _type_descriptor(std::move(type_descriptor)), + _type(std::move(type)), + _name(std::move(name)), + _record_reader(std::move(record_reader)) {} + + int file_column_id() const override { return _file_column_id; } + int parquet_column_ordinal() const override { return _parquet_column_ordinal; } + const DataTypePtr& type() const override { return _type; } + const std::string& name() const override { return _name; } + + Status read_batch(int64_t batch_rows, MutableColumnPtr* result_column, + int64_t* rows_read) override; + Status skip(int64_t rows) override; + Status read_selected(const std::vector<uint16_t>& selection, uint16_t selected_rows, + int64_t batch_rows, MutableColumnPtr* result_column) override; + + const ::parquet::ColumnDescriptor* descriptor() const { return _descriptor; } + const std::shared_ptr<::parquet::internal::RecordReader>& record_reader() const { + return _record_reader; + } + +private: + int _file_column_id = -1; + int _parquet_column_ordinal = -1; + const ::parquet::ColumnDescriptor* _descriptor = nullptr; + ParquetTypeDescriptor _type_descriptor; + DataTypePtr _type; + std::string _name; + std::shared_ptr<::parquet::internal::RecordReader> _record_reader; +}; + +class StructColumnReader final : public ParquetColumnReader { +public: + StructColumnReader(const ParquetColumnSchema& schema, + std::vector<std::unique_ptr<ParquetColumnReader>> children) + : _field_id(schema.field_id), + _type(schema.type), + _name(schema.name), + _children(std::move(children)) {} + + int file_column_id() const override { return _field_id; } + int parquet_column_ordinal() const override { return -1; } + const DataTypePtr& type() const override { return _type; } + const std::string& name() const override { return _name; } + + Status read_batch(int64_t batch_rows, MutableColumnPtr* result_column, + int64_t* rows_read) override; + Status skip(int64_t rows) override; + +private: + int _field_id = -1; + DataTypePtr _type; + std::string _name; + std::vector<std::unique_ptr<ParquetColumnReader>> _children; +}; + +template <typename InsertValue> +Status append_record_reader_values(const PrimitiveColumnReader& column_reader, + ::parquet::internal::RecordReader& record_reader, + int64_t records_read, MutableColumnPtr* result_column, + InsertValue&& insert_value) { + auto column = column_reader.type()->create_column(); + const auto* values = record_reader.values(); + if (values == nullptr && record_reader.values_written() > 0) { + return Status::Corruption("Parquet record reader returned null values buffer for column {}", + column_reader.name()); + } + + if (column_reader.descriptor()->max_definition_level() == 0) { + if (record_reader.values_written() != records_read) { + return Status::Corruption( + "Invalid required parquet record read result for column {}: values={}, " + "records={}", + column_reader.name(), record_reader.values_written(), records_read); + } + for (int64_t value_idx = 0; value_idx < records_read; ++value_idx) { + RETURN_IF_ERROR(insert_value(*column, values, value_idx)); + } + } else { + const int16_t max_definition_level = column_reader.descriptor()->max_definition_level(); + auto* def_levels = record_reader.def_levels(); + if (def_levels == nullptr && records_read > 0) { + return Status::Corruption( + "Parquet record reader returned null definition levels for nullable column {}", + column_reader.name()); + } + if (record_reader.read_dense_for_nullable()) { + return Status::NotSupported( + "Dense nullable parquet record reader is not supported for column {}", + column_reader.name()); + } + if (record_reader.values_written() != records_read) { + return Status::Corruption( + "Invalid nullable parquet record read result for column {}: values={}, " + "records={}", + column_reader.name(), record_reader.values_written(), records_read); + } + for (int64_t record_idx = 0; record_idx < records_read; ++record_idx) { + if (def_levels[record_idx] == max_definition_level) { + RETURN_IF_ERROR(insert_value(*column, values, record_idx)); + } else { + column->insert_data(nullptr, 0); + } + } + } + + *result_column = std::move(column); + return Status::OK(); +} + +Status read_records(PrimitiveColumnReader& column_reader, int64_t batch_rows, + ::parquet::internal::RecordReader** record_reader, int64_t* rows_read) { + auto reader = column_reader.record_reader(); + if (reader == nullptr) { + return Status::InternalError("Parquet record reader is not initialized for column {}", + column_reader.name()); + } + + int64_t records_read = 0; + try { + reader->Reset(); + reader->Reserve(batch_rows); + records_read = reader->ReadRecords(batch_rows); + } catch (const ::parquet::ParquetException& e) { + return Status::Corruption("Failed to read parquet records for column {}: {}", + column_reader.name(), e.what()); + } catch (const std::exception& e) { + return Status::InternalError("Failed to read parquet records for column {}: {}", + column_reader.name(), e.what()); + } + if (records_read < 0 || records_read > batch_rows) { + return Status::Corruption("Invalid parquet record read result for column {}: {}", + column_reader.name(), records_read); + } + *record_reader = reader.get(); + *rows_read = records_read; + return Status::OK(); +} + +template <PrimitiveType DorisType, typename ParquetValueType> +Status read_record_physical_column(PrimitiveColumnReader& column_reader, int64_t batch_rows, + MutableColumnPtr* result_column, int64_t* rows_read) { + ::parquet::internal::RecordReader* record_reader = nullptr; + RETURN_IF_ERROR(read_records(column_reader, batch_rows, &record_reader, rows_read)); + using DorisCppType = typename PrimitiveTypeTraits<DorisType>::CppType; + return append_record_reader_values( + column_reader, *record_reader, *rows_read, result_column, + [](IColumn& column, const uint8_t* values, int64_t value_idx) { + const auto* typed_values = reinterpret_cast<const ParquetValueType*>(values); + DorisCppType doris_value = static_cast<DorisCppType>(typed_values[value_idx]); + column.insert_data(reinterpret_cast<const char*>(&doris_value), + sizeof(DorisCppType)); + return Status::OK(); + }); +} + +template <typename ParquetValueType, typename InsertValue> +Status read_record_value_column(PrimitiveColumnReader& column_reader, int64_t batch_rows, + MutableColumnPtr* result_column, int64_t* rows_read, + InsertValue&& insert_value) { + ::parquet::internal::RecordReader* record_reader = nullptr; + RETURN_IF_ERROR(read_records(column_reader, batch_rows, &record_reader, rows_read)); + return append_record_reader_values( + column_reader, *record_reader, *rows_read, result_column, + [&](IColumn& column, const uint8_t* values, int64_t value_idx) { + const auto* typed_values = reinterpret_cast<const ParquetValueType*>(values); + return insert_value(column, typed_values[value_idx]); + }); +} + +Status get_binary_chunks(const PrimitiveColumnReader& column_reader, + ::parquet::internal::RecordReader& record_reader, + std::vector<std::shared_ptr<::arrow::Array>>* chunks) { + auto* binary_reader = dynamic_cast<::parquet::internal::BinaryRecordReader*>(&record_reader); + if (binary_reader == nullptr) { + return Status::InternalError("Parquet binary record reader is not available for column {}", + column_reader.name()); + } + *chunks = binary_reader->GetBuilderChunks(); + return Status::OK(); +} + +template <typename InsertValue> +Status append_binary_record_reader_values(const PrimitiveColumnReader& column_reader, + ::parquet::internal::RecordReader& record_reader, + int64_t records_read, MutableColumnPtr* result_column, + InsertValue&& insert_value) { + std::vector<std::shared_ptr<::arrow::Array>> chunks; + RETURN_IF_ERROR(get_binary_chunks(column_reader, record_reader, &chunks)); + auto column = column_reader.type()->create_column(); + int64_t appended_rows = 0; + for (const auto& chunk : chunks) { + if (chunk == nullptr) { + return Status::Corruption( + "Parquet binary record reader returned null chunk for column {}", + column_reader.name()); + } + for (int64_t row_idx = 0; row_idx < chunk->length(); ++row_idx) { + if (chunk->IsNull(row_idx)) { + column->insert_data(nullptr, 0); + } else { + RETURN_IF_ERROR(insert_value(*column, *chunk, row_idx)); + } + ++appended_rows; + } + } + if (appended_rows != records_read) { + return Status::Corruption( + "Invalid parquet binary record read result for column {}: rows={}, records={}", + column_reader.name(), appended_rows, records_read); + } + *result_column = std::move(column); + return Status::OK(); +} + +struct RowRange { Review Comment: 和segment_v2里的RowRange统一起来吧 -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
