SteNicholas commented on code in PR #222: URL: https://github.com/apache/paimon-cpp/pull/222#discussion_r3840610347
########## src/paimon/core/table/format/format_table_read.cpp: ########## @@ -0,0 +1,438 @@ +/* + * 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/table/format/format_table_read.h" + +#include <algorithm> +#include <map> +#include <set> +#include <string> +#include <utility> +#include <vector> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/predicate/predicate_validator.h" +#include "paimon/common/reader/concat_batch_reader.h" +#include "paimon/common/reader/predicate_batch_reader.h" +#include "paimon/common/utils/arrow/mem_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/options_utils.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/core/table/format/format_path_validation.h" +#include "paimon/core/table/format/lazy_concat_batch_reader.h" +#include "paimon/core/table/format/limit_batch_reader.h" +#include "paimon/core/table/format/partition_completing_batch_reader.h" +#include "paimon/defs.h" +#include "paimon/format/file_format.h" +#include "paimon/format/file_format_factory.h" +#include "paimon/fs/file_system.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/predicate/predicate_utils.h" +#include "paimon/table/format/format_data_split.h" + +namespace paimon { + +namespace { + +/// Whether every field the predicate names is one the reader will produce. +Result<bool> PredicateFitsSchema(const std::shared_ptr<Predicate>& predicate, + const std::set<std::string>& available_fields) { + std::set<std::string> named_fields; + PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate, &named_fields)); + for (const std::string& field : named_fields) { + if (available_fields.find(field) == available_fields.end()) { + return false; + } + } + return true; +} + +std::set<std::string> FieldNamesOf(const arrow::FieldVector& fields) { + std::set<std::string> names; + for (const std::shared_ptr<arrow::Field>& field : fields) { + names.insert(field->name()); + } + return names; +} + +} // namespace + +/// Everything the read needs that the public header should not have to name. +class FormatTableRead::Impl { + public: + std::shared_ptr<FormatTable> table; + /// Columns the reader returns, in the order it returns them. + std::shared_ptr<arrow::Schema> read_schema; + /// The read columns that live in the data files, in table order. + std::shared_ptr<arrow::Schema> data_read_schema; + /// For each read column, its index in `data_read_schema`, or -1 for a partition column. + std::vector<int32_t> read_to_data_index; + std::optional<int32_t> limit; + /// The conjuncts the file readers can evaluate, i.e. those naming only columns the file is + /// asked for. Null when there is none. + std::shared_ptr<Predicate> pushdown_predicate; + /// The predicate the returned reader applies exactly, or null when the caller filters itself. + std::shared_ptr<Predicate> filter_predicate; + std::shared_ptr<MemoryPool> pool; + std::string format_identifier; + int32_t batch_size = 1024; +}; + +FormatTableRead::FormatTableRead(std::unique_ptr<Impl> impl, + const std::shared_ptr<MemoryPool>& pool) + : TableRead(pool), impl_(std::move(impl)) {} + +FormatTableRead::~FormatTableRead() = default; + +Result<std::unique_ptr<FormatTableRead>> FormatTableRead::Create( + const std::shared_ptr<FormatTable>& table, + const std::optional<std::vector<std::string>>& projection, const std::optional<int32_t>& limit, + const std::shared_ptr<MemoryPool>& pool, const std::shared_ptr<Predicate>& predicate, + bool enable_predicate_filter) { + if (table == nullptr) { + return Status::Invalid("format table read requires a table"); + } + std::shared_ptr<MemoryPool> memory_pool = pool != nullptr ? pool : GetDefaultPool(); + + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> c_schema, table->GetArrowSchema()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema> table_schema, + arrow::ImportSchema(c_schema.get())); + + const std::vector<std::string>& partition_keys = table->PartitionKeys(); + auto is_partition_key = [&partition_keys](const std::string& name) { + return std::find(partition_keys.begin(), partition_keys.end(), name) != + partition_keys.end(); + }; + + arrow::FieldVector read_fields; + if (projection) { + read_fields.reserve(projection->size()); + std::set<std::string> projected; + for (const std::string& name : *projection) { + // A read column is looked up by name, so the same column twice has no meaning to + // act on. + if (!projected.insert(name).second) { + return Status::Invalid(fmt::format( + "column '{}' appears more than once in the projection, which paimon-cpp does " + "not allow", + name)); + } + std::shared_ptr<arrow::Field> field = table_schema->GetFieldByName(name); + if (field == nullptr) { + return Status::Invalid( + fmt::format("field '{}' is not a column of table {}", name, table->FullName())); + } + read_fields.push_back(std::move(field)); + } + } else { + read_fields = table_schema->fields(); + } + if (read_fields.empty()) { + return Status::Invalid("format table read requires at least one column to read"); + } + + auto impl = std::make_unique<Impl>(); + impl->table = table; + impl->read_schema = arrow::schema(read_fields); + impl->limit = limit; + impl->pool = memory_pool; + impl->format_identifier = FormatTable::FormatToString(table->GetFormat()); + + // Every non-partition column of the table, in table order: what a data file holds. + arrow::FieldVector file_fields; + for (const std::shared_ptr<arrow::Field>& field : table_schema->fields()) { + if (!is_partition_key(field->name())) { + file_fields.push_back(field); + } + } + if (file_fields.empty()) { + return Status::Invalid( + fmt::format("format table {} has no non-partition column, so its files hold nothing to " + "read", + table->FullName())); + } + + // Fields the predicate names, checked against the table up front: nothing downstream could + // evaluate a field the table does not have. + std::set<std::string> predicate_fields; + if (predicate != nullptr) { + PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate, &predicate_fields)); + for (const std::string& field : predicate_fields) { + if (table_schema->GetFieldIndex(field) < 0) { + return Status::Invalid(fmt::format( + "predicate field '{}' is not a column of table {}", field, table->FullName())); + } + } + // The rest is what every read path in this library checks: that each conjunct's declared + // type is the column's, and that its literals are neither null nor of another type. The + // field index is not checked, since everything downstream resolves a field by name. + PAIMON_RETURN_NOT_OK(PredicateValidator::ValidatePredicateWithSchema( + *table_schema, predicate, /*validate_field_idx=*/false)); + PAIMON_RETURN_NOT_OK(PredicateValidator::ValidatePredicateWithLiterals(predicate)); + } + + // The columns read out of a file, in the order the caller asked for them. + arrow::FieldVector data_read_fields; + for (const std::shared_ptr<arrow::Field>& read_field : read_fields) { + if (!is_partition_key(read_field->name())) { + data_read_fields.push_back(read_field); + } + } + // A column only the predicate names is read as well, so the format can skip what its own + // statistics let it skip; the batch drops it again. + for (const std::shared_ptr<arrow::Field>& field : file_fields) { + if (predicate_fields.find(field->name()) == predicate_fields.end()) { + continue; + } + bool already_read = std::any_of(data_read_fields.begin(), data_read_fields.end(), + [&field](const std::shared_ptr<arrow::Field>& read_field) { + return read_field->name() == field->name(); + }); + if (!already_read) { + data_read_fields.push_back(field); + } + } + if (data_read_fields.empty()) { + // A partition-only projection still has to read something: only the file says how many + // rows to repeat the constant values for. + data_read_fields.push_back(file_fields.front()); + } + impl->data_read_schema = arrow::schema(data_read_fields); + + impl->read_to_data_index.reserve(read_fields.size()); + for (const std::shared_ptr<arrow::Field>& read_field : read_fields) { + if (is_partition_key(read_field->name())) { + impl->read_to_data_index.push_back(-1); + } else { + impl->read_to_data_index.push_back( + static_cast<int32_t>(impl->data_read_schema->GetFieldIndex(read_field->name()))); + } + } + + if (predicate != nullptr) { + // A conjunct goes to whichever can evaluate every field it names: the file reader, which + // sees the non-partition columns, or the batch, which carries the projected ones plus the + // partitions. That leaves exactly the conjuncts naming a partition key out of the + // pushdown. + const std::set<std::string> data_read_field_names = + FieldNamesOf(impl->data_read_schema->fields()); + const std::set<std::string> read_field_names = FieldNamesOf(read_fields); + + std::vector<std::shared_ptr<Predicate>> pushdown_conjuncts; + std::vector<std::shared_ptr<Predicate>> filter_conjuncts; + for (const std::shared_ptr<Predicate>& conjunct : PredicateUtils::SplitAnd(predicate)) { + PAIMON_ASSIGN_OR_RAISE(bool fits_file, + PredicateFitsSchema(conjunct, data_read_field_names)); + if (fits_file) { + pushdown_conjuncts.push_back(conjunct); + } + PAIMON_ASSIGN_OR_RAISE(bool fits_batch, + PredicateFitsSchema(conjunct, read_field_names)); + if (fits_batch) { + filter_conjuncts.push_back(conjunct); + } else if (enable_predicate_filter) { + // Dropping it would return rows the caller asked to have filtered out. + return Status::Invalid(fmt::format( + "predicate '{}' of table {} names a column the projection drops, so the " + "reader cannot apply it; project that column or leave " + "enable_predicate_filter off", + conjunct->ToString(), table->FullName())); + } + } + if (!pushdown_conjuncts.empty()) { + PAIMON_ASSIGN_OR_RAISE(impl->pushdown_predicate, + PredicateBuilder::And(pushdown_conjuncts)); + } + if (enable_predicate_filter && !filter_conjuncts.empty()) { + PAIMON_ASSIGN_OR_RAISE(impl->filter_predicate, PredicateBuilder::And(filter_conjuncts)); + } + } + + PAIMON_ASSIGN_OR_RAISE(impl->batch_size, OptionsUtils::GetValueFromMap<int32_t>( + table->Options(), Options::READ_BATCH_SIZE, 1024)); + if (impl->batch_size <= 0) { + return Status::Invalid(fmt::format("{} must be larger than 0, but is {}", + Options::READ_BATCH_SIZE, impl->batch_size)); + } + + return std::unique_ptr<FormatTableRead>(new FormatTableRead(std::move(impl), memory_pool)); +} + +Result<std::unique_ptr<BatchReader>> FormatTableRead::CreateSplitReader( + const std::shared_ptr<Split>& split) { + auto format_split = std::dynamic_pointer_cast<FormatDataSplit>(split); + if (format_split == nullptr) { + return Status::Invalid("format table read only accepts a FormatDataSplit"); + } + + // A split is public and can be decoded from bytes this process did not write. Deserialization + // can only check the numbers; whether a file belongs to this table is asked here. + PAIMON_RETURN_NOT_OK(ValidatePartitionKeys(impl_->table, format_split->Partition(), "split")); + for (const FormatDataSplit::FileMeta& file : format_split->Files()) { + PAIMON_RETURN_NOT_OK( + ValidatePathUnderLocation(file.file_path, impl_->table->Location(), "split")); + // A split mixing partitions would read rows back under values they never had. + PAIMON_RETURN_NOT_OK(ValidateFileInPartition(impl_->table, file.file_path, + format_split->Partition(), "split")); + PAIMON_RETURN_NOT_OK(ValidateFileIsVisible(impl_->table, file.file_path, "split")); + if (file.file_size < 0) { + return Status::Invalid(fmt::format("split gives {} a negative size", file.file_path)); + } + } + + std::vector<PartitionCompletingBatchReader::ColumnSource> column_sources; + column_sources.reserve(impl_->read_to_data_index.size()); + const std::map<std::string, std::string>& partition = format_split->Partition(); + for (int32_t i = 0; i < impl_->read_schema->num_fields(); i++) { + const std::shared_ptr<arrow::Field>& read_field = impl_->read_schema->field(i); + int32_t data_index = impl_->read_to_data_index[i]; + if (data_index >= 0) { + column_sources.push_back({data_index, std::string(), false}); + continue; + } + auto partition_iter = partition.find(read_field->name()); + if (partition_iter == partition.end()) { + return Status::Invalid(fmt::format( + "split does not carry a value for partition field '{}'", read_field->name())); + } + // A directory named after the default partition name stands for a null partition value. + bool is_null = partition_iter->second == impl_->table->PartitionDefaultName(); + column_sources.push_back({-1, partition_iter->second, is_null}); + } + + // The format and builder depend only on the table, so one serves the whole split. Shared, + // because each file's reader is built later, when that file is reached. + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<FileFormat> owned_format, + FileFormatFactory::Get(impl_->format_identifier, impl_->table->Options())); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReaderBuilder> owned_builder, + owned_format->CreateReaderBuilder(impl_->batch_size)); + owned_builder->WithMemoryPool(impl_->pool); + std::shared_ptr<FileFormat> file_format(std::move(owned_format)); + std::shared_ptr<ReaderBuilder> reader_builder(std::move(owned_builder)); + + // Captured by value, so a file's reader stays valid however long this `FormatTableRead` + // lives. + std::shared_ptr<FormatTable> table = impl_->table; + std::shared_ptr<arrow::Schema> data_read_schema = impl_->data_read_schema; + std::shared_ptr<arrow::Schema> read_schema = impl_->read_schema; + std::shared_ptr<Predicate> pushdown_predicate = impl_->pushdown_predicate; + std::shared_ptr<MemoryPool> pool = impl_->pool; + auto shared_column_sources = + std::make_shared<std::vector<PartitionCompletingBatchReader::ColumnSource>>( + std::move(column_sources)); + + // Each file is named alongside the factory that opens it, so every failure says which file + // it was. + std::vector<LazyConcatBatchReader::Source> sources; + sources.reserve(format_split->Files().size()); + for (const FormatDataSplit::FileMeta& file : format_split->Files()) { + LazyConcatBatchReader::Source source; + source.name = file.file_path; + source.open = [table, file_format, reader_builder, data_read_schema, read_schema, + shared_column_sources, pool, pushdown_predicate, + file]() -> Result<std::unique_ptr<BatchReader>> { + // Captured only to outlive the builder it created, which some formats need. + (void)file_format; + // `Open` takes a status the caller vouches for, and the split's size is whatever it + // was given: a stale or forged length would truncate an object-store read or send it + // past the end. So the file system is asked for the real one. + PAIMON_ASSIGN_OR_RAISE(FileStatus status, + table->GetFileSystem()->GetFileStatus(file.file_path)); + if (status.IsDir()) { + return Status::Invalid("the split names a directory, not a data file"); + } + if (file.file_size != status.GetLen()) { + return Status::Invalid(fmt::format( + "the split says it is {} bytes but it is {}; the plan was made against a " + "different version of the file", + file.file_size, status.GetLen())); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InputStream> input_stream, + table->GetFileSystem()->Open(status)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileBatchReader> file_reader, + reader_builder->Build(input_stream)); + + ::ArrowSchema c_read_schema; + ArrowSchemaMarkReleased(&c_read_schema); + ScopeGuard read_schema_guard( + [&c_read_schema]() { ArrowSchemaRelease(&c_read_schema); }); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*data_read_schema, &c_read_schema)); + PAIMON_RETURN_NOT_OK(file_reader->SetReadSchema(&c_read_schema, pushdown_predicate, + /*selection_bitmap=*/std::nullopt)); + + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr<PartitionCompletingBatchReader> reader, + PartitionCompletingBatchReader::Create(std::move(file_reader), data_read_schema, Review Comment: Done in substance. - `TableScan::Create()`, `TableRead::Create()`, `FileStoreWrite::Create()` and `FileStoreCommit::Create()` recognise a format table from the schema under the table path and dispatch to it, so a `BatchReader` still comes from `TableRead` and a caller keeps one interface. One `FormatTableLoader` answers that question for all four and hands the schema it read to the managed path, so the schema file is read once. - Everything the read needs is taken from `ReadContext`: the columns, the predicate and whether to apply it exactly, the pool and executor, prefetch, the read-ahead cache and its config, and the block cache. What a format table cannot honour is refused by name rather than dropped. - File opening is shared with the managed path through a new `DataFileReaderFactory`, which `AbstractSplitRead` now uses too, so prefetch and caching behave identically on both paths. `RawFileSplitRead` itself is not reused: it is built around `DataFileMeta`, `DataFilePathFactory`, deletion vectors and bitmap indexes, none of which a format table has. `DataFileReaderFactory` is the seam the two paths share instead. ########## src/paimon/core/table/format/format_table_scan.cpp: ########## @@ -0,0 +1,249 @@ +/* + * 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/table/format/format_table_scan.h" + +#include <algorithm> +#include <string> +#include <utility> + +#include "fmt/format.h" +#include "paimon/common/options/memory_size.h" +#include "paimon/common/utils/bin_packing.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/core/table/format/format_file_listing.h" +#include "paimon/core/table/format/format_path_validation.h" +#include "paimon/core/table/source/plan_impl.h" +#include "paimon/core/utils/partition_path_utils.h" +#include "paimon/defs.h" +#include "paimon/fs/file_system.h" +#include "paimon/table/format/format_data_split.h" + +namespace paimon { + +namespace { + +/// Defaults of `source.split.target-size` and `source.split.open-file-cost`. +constexpr int64_t kDefaultTargetSplitSize = 128 * 1024 * 1024; +constexpr int64_t kDefaultOpenFileCost = 4 * 1024 * 1024; + Review Comment: Done. `FormatTableScan`, `FormatTableRead`, `FormatTableWrite` and `FormatTable::Create()` all read through `CoreOptions`, so a default lives in one place. `CoreOptions` gained `FileSuffixIncludeCompression()`, `FormatTablePartitionOnlyValueInPath()`, `MetastorePartitionedTable()` and `FormatTableFileCompression()` (the `file.compression` then `format-table.file.compression` then `compression` chain), with `CoreOptionsTest.TestFormatTableOptions` and `TestFormatTableFileCompression` covering them. ########## src/paimon/core/table/format/limit_batch_reader.cpp: ########## @@ -0,0 +1,83 @@ +/* + * 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/core/table/format/limit_batch_reader.h" + +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "paimon/common/reader/reader_utils.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/status.h" + +namespace paimon { + +Result<std::unique_ptr<LimitBatchReader>> LimitBatchReader::Create( + std::unique_ptr<BatchReader>&& reader, int64_t limit, const std::shared_ptr<MemoryPool>& pool) { + if (reader == nullptr) { + return Status::Invalid("limit reader requires a reader"); + } + if (limit < 0) { + return Status::Invalid("limit must not be negative"); + } + return std::unique_ptr<LimitBatchReader>(new LimitBatchReader(std::move(reader), limit, pool)); +} + +Result<BatchReader::ReadBatch> LimitBatchReader::NextBatch() { + if (returned_rows_ >= limit_) { + return BatchReader::MakeEofBatch(); + } + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, reader_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + return batch; + } + int64_t batch_rows = batch.first->length; + // Compared by subtraction: `returned_rows_` never passes `limit_`, so the difference cannot + // overflow where the sum of two int64 row counts could. + if (batch_rows <= limit_ - returned_rows_) { + returned_rows_ += batch_rows; + return batch; + } + + // Only the batch's first rows are wanted. Slicing keeps the buffers shared, but a sliced + // array carries an offset that `BatchReader` forbids, so the offsets are rebased. + int64_t keep_rows = limit_ - returned_rows_; + // Importing takes the batch over on success only, and a `unique_ptr<ArrowArray>` frees the + // struct without releasing what it points at. + arrow::Result<std::shared_ptr<arrow::RecordBatch>> imported = + arrow::ImportRecordBatch(batch.first.get(), batch.second.get()); + if (!imported.ok()) { + ReaderUtils::ReleaseReadBatch(std::move(batch)); + return ToPaimonStatus(imported.status()); + } Review Comment: Moot: the file is gone with the limit reader. ########## src/paimon/core/table/format/lazy_concat_batch_reader.cpp: ########## @@ -0,0 +1,122 @@ +/* + * 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/core/table/format/lazy_concat_batch_reader.h" + +#include <string> +#include <utility> + +// `ReadBatch` holds `unique_ptr`s to these, so destroying one needs their definitions. +#include "arrow/c/abi.h" +#include "fmt/format.h" +#include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/reader/reader_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" + +namespace paimon { + +namespace { +/// Says which file a failure came from, keeping the status code it carried. +Status WithFileName(const std::string& name, const Status& status) { + return Status(status.code(), fmt::format("cannot read {}: {}", name, status.message())); +} +} // namespace + +LazyConcatBatchReader::LazyConcatBatchReader(std::vector<Source>&& sources, + const std::shared_ptr<MemoryPool>& pool) + : arrow_pool_(GetArrowPool(pool)), + sources_(std::move(sources)), + closed_metrics_(std::make_shared<MetricsImpl>()) {} + +LazyConcatBatchReader::~LazyConcatBatchReader() { + DoClose(); +} + +void LazyConcatBatchReader::CloseCurrent() { + if (current_reader_ == nullptr) { + return; + } + // Taken before the reader goes, or everything read through it is missing from the totals. + std::shared_ptr<Metrics> metrics = current_reader_->GetReaderMetrics(); + current_reader_->Close(); + if (metrics != nullptr) { + closed_metrics_->Merge(metrics); + } + // Kept until this reader goes: a batch it handed out is allocated from a pool it owns, and + // the caller only holds this one. Closing already released the file, so little stays behind. + closed_readers_.push_back(std::move(current_reader_)); + current_name_.clear(); +} + +Result<BatchReader::ReadBatchWithBitmap> LazyConcatBatchReader::NextBatchWithBitmap() { + while (true) { + if (current_reader_ == nullptr) { + if (next_source_ >= sources_.size()) { + return BatchReader::MakeEofBatchWithBitmap(); + } + Source& source = sources_[next_source_++]; + Result<std::unique_ptr<BatchReader>> opened = source.open(); + if (!opened.ok()) { + return WithFileName(source.name, opened.status()); + } + current_reader_ = std::move(opened).value(); + if (current_reader_ == nullptr) { Review Comment: Good catch, fixed. The source index only advances after a successful open, and the reader remembers the failure it stopped at and keeps answering with it, so a second `NextBatch()` never reaches the next file. `LazyConcatBatchReaderTest.TestAFailureIsTerminal` asserts both the repeated error and that the second file was never opened. ########## src/paimon/core/table/format/format_table_write.cpp: ########## @@ -0,0 +1,598 @@ +/* + * 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/table/format/format_table_write.h" + +#include <algorithm> +#include <limits> +#include <map> +#include <optional> +#include <string> +#include <string_view> +#include <utility> + +#include "arrow/api.h" +#include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" +#include "fmt/format.h" +#include "paimon/common/options/memory_size.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/field_type_utils.h" +#include "paimon/common/utils/hadoop_compression.h" +#include "paimon/common/utils/options_utils.h" +#include "paimon/common/utils/path_util.h" +#include "paimon/common/utils/scope_guard.h" +#include "paimon/common/utils/string_utils.h" +#include "paimon/core/casting/cast_executor.h" +#include "paimon/core/casting/cast_executor_factory.h" +#include "paimon/core/table/format/format_file_naming.h" +#include "paimon/core/table/format/format_path_validation.h" +#include "paimon/defs.h" +#include "paimon/format/file_format.h" +#include "paimon/format/file_format_factory.h" +#include "paimon/format/format_writer.h" +#include "paimon/format/writer_builder.h" +#include "paimon/fs/file_system.h" +#include "paimon/logging.h" + +namespace paimon { + +namespace { + +/// Every failure logged below is one this write carried on past; without a trace an abandoned +/// temp file cannot be accounted for afterwards. +Logger* WriteLogger() { + static std::unique_ptr<Logger> logger = Logger::GetLogger("FormatTableWrite"); + return logger.get(); +} + +/// Default target size of a data file, matching the append table default. +constexpr int64_t kDefaultTargetFileSize = 256 * 1024 * 1024; +/// Default target row count of a data file: unbounded, so only the size decides. +constexpr int64_t kDefaultTargetFileRowNum = std::numeric_limits<int64_t>::max(); +/// Rolling is checked between batches, so a batch is the finest granularity a file can roll at. +constexpr int32_t kDefaultWriteBatchSize = 1024; + +/// The extension a compression adds to a data file's name: a hadoop compression by its own +/// extension, anything else by the option's text, no compression by nothing. +std::string CompressionFileExtension(const std::string& compression) { + if (compression.empty()) { + return std::string(); + } + std::optional<HadoopCompression::Kind> kind = HadoopCompression::FromName(compression); + if (kind) { + return HadoopCompression::ToFileExtension(*kind); + } + return compression; +} + +/// Renders a partition column as the text a partition directory is named with, which is also the +/// text the read path parses back out of that name. +Result<std::shared_ptr<arrow::StringArray>> RenderPartitionColumnAsText( + const std::shared_ptr<arrow::Array>& column, const std::string& field_name, + arrow::MemoryPool* pool) { + if (column->type_id() == arrow::Type::STRING) { + return checked_pointer_cast<arrow::StringArray>(column); + } + PAIMON_ASSIGN_OR_RAISE(FieldType source_type, + FieldTypeUtils::ConvertToFieldType(column->type()->id())); + std::shared_ptr<CastExecutor> cast_executor = + CastExecutorFactory::GetCastExecutorFactory()->GetCastExecutor(source_type, + FieldType::STRING); + if (cast_executor == nullptr) { + return Status::NotImplemented( + fmt::format("cannot name a partition directory after field '{}' of type {}", field_name, + column->type()->ToString())); + } + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> casted, + cast_executor->Cast(column, arrow::utf8(), pool)); + return checked_pointer_cast<arrow::StringArray>(casted); +} + +} // namespace + +/// The file currently being written for one partition. +struct FormatTableWriteFile { + std::shared_ptr<OutputStream> out; + std::unique_ptr<FormatWriter> writer; + std::string temp_file_path; + std::string file_path; + int64_t record_count = 0; +}; + +/// Everything the write needs that the public header should not have to name. +class FormatTableWrite::Impl { + public: + /// The directory a partition's files belong in, derived once per partition: `Write()` needs + /// it for every batch, and deriving it escapes and then re-walks a whole path. + Result<std::string> DirectoryOf(const std::map<std::string, std::string>& partition); + + /// Opens a new hidden file in `directory` for the partition that directory stands for. + Result<FormatTableWriteFile> OpenFile(const std::string& directory); + + /// Closes the file open in `directory` and records it for committing. + Status FinishFile(const std::string& directory); + + /// Checks that every row really belongs to the partition the batch declares. Partition + /// columns are not written to the file, so a disagreeing row would be read back with the + /// declared value and its own lost beyond recovery. + Status ValidatePartitionColumns( + const std::shared_ptr<arrow::StructArray>& batch, + const std::vector<std::pair<std::string, std::string>>& ordered_partition); + + std::shared_ptr<FormatTable> table; + std::shared_ptr<MemoryPool> pool; + std::unique_ptr<arrow::MemoryPool> arrow_pool; + /// Full table schema, used to check the incoming batch. + std::shared_ptr<arrow::Schema> table_schema; + std::shared_ptr<arrow::DataType> table_struct_type; + /// Columns actually stored in the files: the table's, minus the partition ones. + std::shared_ptr<arrow::Schema> data_schema; + /// Index in the table schema of each data column, so a batch can be projected without + /// re-deriving the mapping every time. + std::vector<int32_t> data_column_indexes; + /// Index in the table schema of each partition column, in partition key order. + std::vector<int32_t> partition_column_indexes; + std::string format_identifier; + std::string file_compression; + int64_t target_file_size = kDefaultTargetFileSize; + int64_t target_file_row_num = kDefaultTargetFileRowNum; + int32_t write_batch_size = kDefaultWriteBatchSize; + FormatFileNaming naming; + + /// The directory each partition seen so far writes into. See `DirectoryOf()`. + std::map<std::map<std::string, std::string>, std::string> partition_directories; + /// Open file per partition, keyed by the partition's directory. + std::map<std::string, FormatTableWriteFile> open_files; + /// Partition values of each open file, by the same key. + std::map<std::string, std::map<std::string, std::string>> open_partitions; + /// Files written and closed but not yet published. They stay here after `PrepareCommit()` + /// hands out a copy, so that an `Abort()` still knows what to remove. + std::vector<FormatCommitMessage> staged_messages; + bool prepared = false; + bool aborted = false; + + /// Why this write will take no more rows, or null while it still will. The two terminal + /// states share the check but need different work from the caller - commit the messages it + /// holds, or start over - so the refusal says which one it hit. + const char* FinishedReason() const { + if (aborted) { + return "format table write has been aborted"; + } + if (prepared) { + return "format table write has already prepared its commit"; + } + return nullptr; + } +}; + +FormatTableWrite::FormatTableWrite(std::unique_ptr<Impl> impl) : impl_(std::move(impl)) {} + +FormatTableWrite::~FormatTableWrite() { + if (impl_ != nullptr && impl_->FinishedReason() == nullptr) { + // `Abort()` reports cleanup failures through the log and comes back OK today; its status + // is still read, here and at the two other callers, in case that ever changes. + Status status = Abort(); + if (!status.ok()) { + PAIMON_LOG_WARN(WriteLogger(), "Failed to abort an abandoned write of table %s: %s", + impl_->table->FullName().c_str(), status.ToString().c_str()); + } + } +} + +Result<std::unique_ptr<FormatTableWrite>> FormatTableWrite::Create( + const std::shared_ptr<FormatTable>& table, const std::shared_ptr<MemoryPool>& pool) { + if (table == nullptr) { + return Status::Invalid("format table write requires a table"); + } + auto impl = std::make_unique<Impl>(); + impl->table = table; + impl->pool = pool != nullptr ? pool : GetDefaultPool(); + + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> c_schema, table->GetArrowSchema()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(impl->table_schema, arrow::ImportSchema(c_schema.get())); + impl->table_struct_type = arrow::struct_(impl->table_schema->fields()); + + const std::vector<std::string>& partition_keys = table->PartitionKeys(); + arrow::FieldVector data_fields; + for (int32_t i = 0; i < impl->table_schema->num_fields(); i++) { + const std::shared_ptr<arrow::Field>& field = impl->table_schema->field(i); + if (std::find(partition_keys.begin(), partition_keys.end(), field->name()) == + partition_keys.end()) { + data_fields.push_back(field); + impl->data_column_indexes.push_back(i); + } + } + if (data_fields.empty()) { + return Status::Invalid(fmt::format( + "format table {} has no non-partition column, so its files would hold nothing", + table->FullName())); + } + impl->data_schema = arrow::schema(data_fields); + + // In partition key order, which is the order the directories nest in. + for (const std::string& partition_key : partition_keys) { + int32_t index = impl->table_schema->GetFieldIndex(partition_key); + if (index < 0) { + return Status::Invalid(fmt::format("partition field '{}' is not a column of table {}", + partition_key, table->FullName())); + } + impl->partition_column_indexes.push_back(index); + } + + impl->format_identifier = FormatTable::FormatToString(table->GetFormat()); + impl->file_compression = table->FileCompression(); + + const std::map<std::string, std::string>& options = table->Options(); + + // parquet and orc record their compression inside the file, so the name keeps the plain + // `.parquet` unless `file.suffix.include.compression` asks for it, and then the compression + // goes in front: `data-<uuid>-0.snappy.parquet`. + const std::string compression_extension = CompressionFileExtension(impl->file_compression); + std::string extension = impl->format_identifier; + if (!compression_extension.empty()) { + PAIMON_ASSIGN_OR_RAISE(bool suffix_include_compression, + OptionsUtils::GetValueFromMap<bool>( + options, Options::FILE_SUFFIX_INCLUDE_COMPRESSION, false)); + if (suffix_include_compression) { + extension = compression_extension + "." + extension; + } + } + + // Read one option at a time rather than through CoreOptions, which would also resolve the + // manifest format a format table has no manifests for. `target-file-size` is a memory size, + // so `"256 mb"` must mean here what it means elsewhere. + auto target_file_size_iter = options.find(Options::TARGET_FILE_SIZE); + if (target_file_size_iter != options.end()) { + PAIMON_ASSIGN_OR_RAISE(impl->target_file_size, + MemorySize::ParseBytes(target_file_size_iter->second)); + } + if (impl->target_file_size <= 0) { + return Status::Invalid(fmt::format("{} must be larger than 0, but is {}", + Options::TARGET_FILE_SIZE, impl->target_file_size)); + } + PAIMON_ASSIGN_OR_RAISE(impl->target_file_row_num, + OptionsUtils::GetValueFromMap<int64_t>( + options, Options::TARGET_FILE_ROW_NUM, kDefaultTargetFileRowNum)); + if (impl->target_file_row_num <= 0) { + return Status::Invalid(fmt::format("{} must be larger than 0, but is {}", + Options::TARGET_FILE_ROW_NUM, + impl->target_file_row_num)); + } + PAIMON_ASSIGN_OR_RAISE(impl->write_batch_size, + OptionsUtils::GetValueFromMap<int32_t>( + options, Options::WRITE_BATCH_SIZE, kDefaultWriteBatchSize)); + if (impl->write_batch_size <= 0) { + return Status::Invalid(fmt::format("{} must be larger than 0, but is {}", + Options::WRITE_BATCH_SIZE, impl->write_batch_size)); + } Review Comment: Done, through `CoreOptions` as well. -- 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]
