feat: add FastAppend

Author: zhjwpku

src/iceberg/update/fast_append.cc

@@ -0,0 +1,340 @@
+/*
+ * 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 "iceberg/update/fast_append.h"
+
+#include <format>
+#include <vector>
+
+#include "iceberg/constants.h"
+#include "iceberg/file_io.h"
+#include "iceberg/manifest/manifest_entry.h"
+#include "iceberg/manifest/manifest_reader.h"
+#include "iceberg/manifest/manifest_writer.h"
+#include "iceberg/snapshot.h"
+#include "iceberg/table_metadata.h"
+#include "iceberg/transaction.h"
+#include "iceberg/util/macros.h"
+
+namespace iceberg {
+
+Result<std::unique_ptr<FastAppend>> FastAppend::Make(
+    const std::string& table_name, std::shared_ptr<Transaction> transaction) {
+  return std::make_unique<FastAppend>(table_name, std::move(transaction));
+}
+
+FastAppend::FastAppend(const std::string& table_name,
+                       std::shared_ptr<Transaction> transaction)
+    : SnapshotUpdate(std::move(transaction)), table_name_(table_name) {}
+
+FastAppend& FastAppend::AppendFile(std::shared_ptr<DataFile> file) {
+  if (file == nullptr) {
+    AddError(ErrorKind::kInvalidArgument, "Invalid data file: null");
+    return *this;
+  }
+
+  if (!file->partition_spec_id.has_value()) {
+    AddError(ErrorKind::kInvalidArgument, "Data file must have partition spec ID");
+    return *this;
+  }
+
+  int32_t spec_id = file->partition_spec_id.value();
+  auto spec_result = Spec(spec_id);
+  if (!spec_result.has_value()) {
+    AddError(spec_result.error());
+    return *this;
+  }
+
+  auto spec = spec_result.value();
+  if (spec == nullptr) {
+    AddError(ErrorKind::kInvalidArgument,
+             "Cannot find partition spec {} for data file: {}", spec_id, file->file_path);
+    return *this;
+  }
+
+  auto& data_files = new_data_files_by_spec_[spec_id];
+  // Check if file already exists (simple check by path)
+  bool is_new = true;
+  for (const auto& existing_file : data_files) {
+    if (existing_file->file_path == file->file_path) {
+      is_new = false;
+      break;
+    }
+  }
+
+  if (is_new) {
+    data_files.push_back(file);
+    has_new_files_ = true;
+    AddFileToSummary(spec, file);
+  }
+
+  return *this;
+}
+
+FastAppend& FastAppend::AppendManifest(const ManifestFile& manifest) {
+  if (manifest.existing_files_count.has_value() && manifest.existing_files_count != 0) {
+    AddError(ErrorKind::kInvalidArgument, "Cannot append manifest with existing files");
+    return *this;
+  }
+  if (manifest.deleted_files_count.has_value() && manifest.deleted_files_count != 0) {
+    AddError(ErrorKind::kInvalidArgument, "Cannot append manifest with deleted files");
+    return *this;
+  }
+  if (manifest.added_snapshot_id != kInvalidSnapshotId &&
+      manifest.added_snapshot_id != -1) {
+    AddError(ErrorKind::kInvalidArgument, "Snapshot id must be assigned during commit");
+    return *this;
+  }
+  if (manifest.sequence_number != -1 &&
+      manifest.sequence_number != TableMetadata::kInitialSequenceNumber) {
+    AddError(ErrorKind::kInvalidArgument,
+             "Sequence number must be assigned during commit");
+    return *this;
+  }
+
+  if (can_inherit_snapshot_id() && (manifest.added_snapshot_id == kInvalidSnapshotId ||
+                                    manifest.added_snapshot_id == -1)) {
+    AddManifestToSummary(manifest);
+    append_manifests_.push_back(manifest);
+  } else {
+    // The manifest must be rewritten with this update's snapshot ID
+    auto copied_manifest_result = CopyManifest(manifest);
+    if (!copied_manifest_result.has_value()) {
+      AddError(copied_manifest_result.error());
+      return *this;
+    }
+    rewritten_append_manifests_.push_back(copied_manifest_result.value());
+  }
+
+  return *this;
+}
+
+FastAppend& FastAppend::ToBranch(const std::string& branch) {
+  auto status = SetTargetBranch(branch);
+  if (!status.has_value()) {
+    AddError(status.error());
+  }
+  return *this;
+}
+
+FastAppend& FastAppend::Set(const std::string& property, const std::string& value) {
+  summary_[property] = value;
+  return *this;
+}
+
+std::string FastAppend::operation() { return DataOperation::kAppend; }
+
+Result<std::vector<ManifestFile>> FastAppend::Apply(
+    const TableMetadata& metadata_to_update, const std::shared_ptr<Snapshot>& snapshot) {
+  std::vector<ManifestFile> manifests;
+
+  ICEBERG_ASSIGN_OR_RAISE(auto new_written_manifests, WriteNewManifests());
+  if (!new_written_manifests.empty()) {
+    manifests.insert(manifests.end(), new_written_manifests.begin(),
+                     new_written_manifests.end());
+  }
+
+  // Transform append manifests and rewritten append manifests with snapshot ID
+  int64_t snapshot_id = SnapshotId();
+  for (const auto& manifest : append_manifests_) {
+    ManifestFile updated = manifest;
+    updated.added_snapshot_id = snapshot_id;
+    manifests.push_back(updated);
+  }
+
+  for (const auto& manifest : rewritten_append_manifests_) {
+    ManifestFile updated = manifest;
+    updated.added_snapshot_id = snapshot_id;
+    manifests.push_back(updated);
+  }
+
+  // Add all manifests from the snapshot
+  if (snapshot != nullptr) {
+    // Use SnapshotCache to get manifests, similar to snapshot_update.cc
+    auto cached_snapshot = SnapshotCache(snapshot.get());
+    auto snapshot_manifests_result =
+        cached_snapshot.Manifests(transaction_->table()->io());
+    if (!snapshot_manifests_result.has_value()) {
+      return snapshot_manifests_result.error();
+    }
+    auto snapshot_manifests_span = snapshot_manifests_result.value();
+    std::vector<ManifestFile> snapshot_manifests(snapshot_manifests_span.begin(),
+                                                 snapshot_manifests_span.end());
+    manifests.insert(manifests.end(), snapshot_manifests.begin(),
+                     snapshot_manifests.end());
+  }
+
+  return manifests;
+}
+
+std::unordered_map<std::string, std::string> FastAppend::Summary() {
+  // Set partition summary limit from table properties
+  // Note: Partition summary limit is handled during manifest writing, not here
+  // The limit affects how many partition summaries are included in manifest metadata
+
+  // Note: In Java, the summary builder has a method to set partition summary limit.
+  // Here we just return the summary map. The partition summary limit is typically
+  // used when building partition summaries, which happens during manifest writing.
+  // We'll handle it in the manifest writing phase if needed.
+
+  return summary_;
+}
+
+void FastAppend::CleanUncommitted(const std::unordered_set<std::string>& committed) {
+  // Clean up new manifests that were written but not committed
+  if (!new_manifests_.empty()) {
+    for (const auto& manifest : new_manifests_) {
+      if (committed.find(manifest.manifest_path) == committed.end()) {
+        std::ignore = transaction_->table()->io()->DeleteFile(manifest.manifest_path);
+      }
+    }
+    new_manifests_.clear();
+  }
+
+  // Clean up only rewritten_append_manifests as they are always owned by the table
+  // Don't clean up append_manifests as they are added to the manifest list and are
+  // not compacted
+  if (!rewritten_append_manifests_.empty()) {
+    for (const auto& manifest : rewritten_append_manifests_) {
+      if (committed.find(manifest.manifest_path) == committed.end()) {
+        std::ignore = transaction_->table()->io()->DeleteFile(manifest.manifest_path);
+      }
+    }
+  }
+}
+
+bool FastAppend::CleanupAfterCommit() const {
+  // Cleanup after committing is disabled for FastAppend unless there are
+  // rewritten_append_manifests because:
+  // 1.) Appended manifests are never rewritten
+  // 2.) Manifests which are written out as part of appendFile are already cleaned
+  //     up between commit attempts in writeNewManifests
+  return !rewritten_append_manifests_.empty();
+}
+
+Result<std::shared_ptr<PartitionSpec>> FastAppend::Spec(int32_t spec_id) {
+  return base().PartitionSpecById(spec_id);
+}
+
+Result<ManifestFile> FastAppend::CopyManifest(const ManifestFile& manifest) {
+  const TableMetadata& current = base();
+  ICEBERG_ASSIGN_OR_RAISE(auto schema, current.Schema());
+  ICEBERG_ASSIGN_OR_RAISE(auto spec,
+                          current.PartitionSpecById(manifest.partition_spec_id));
+
+  // Read the manifest entries
+  ICEBERG_ASSIGN_OR_RAISE(
+      auto reader,
+      ManifestReader::Make(manifest, transaction_->table()->io(), schema, spec));
+  ICEBERG_ASSIGN_OR_RAISE(auto entries, reader->Entries());
+
+  // Create a new manifest writer
+  // Generate a unique manifest path using the transaction's metadata location
+  std::string filename = std::format("copy-m{}.avro", copy_manifest_count_++);
+  std::string new_manifest_path = transaction_->MetadataFileLocation(filename);
+  int64_t snapshot_id = SnapshotId();
+  ICEBERG_ASSIGN_OR_RAISE(
+      auto writer, ManifestWriter::MakeWriter(
+                       current.format_version, snapshot_id, new_manifest_path,
+                       transaction_->table()->io(), spec, schema, ManifestContent::kData,
+                       /*first_row_id=*/current.next_row_id));
+
+  // Write all entries as added entries with the new snapshot ID
+  for (auto& entry : entries) {
+    ICEBERG_PRECHECK(entry.status == ManifestStatus::kAdded,
+                     "Manifest to copy must only contain added entries");
+    entry.snapshot_id = snapshot_id;
+    ICEBERG_RETURN_UNEXPECTED(writer->WriteAddedEntry(entry));
+  }
+
+  ICEBERG_RETURN_UNEXPECTED(writer->Close());
+  ICEBERG_ASSIGN_OR_RAISE(auto new_manifest, writer->ToManifestFile());
+
+  // Update summary with the copied manifest
+  AddManifestToSummary(new_manifest);
+
+  return new_manifest;
+}
+
+Result<std::vector<ManifestFile>> FastAppend::WriteNewManifests() {
+  // If there are new files and manifests were already written, clean them up
+  if (has_new_files_ && !new_manifests_.empty()) {
+    for (const auto& manifest : new_manifests_) {
+      std::ignore = transaction_->table()->io()->DeleteFile(manifest.manifest_path);
+    }
+    new_manifests_.clear();
+  }
+
+  // Write new manifests if there are new data files
+  if (new_manifests_.empty() && !new_data_files_by_spec_.empty()) {
+    for (const auto& [spec_id, data_files] : new_data_files_by_spec_) {
+      ICEBERG_ASSIGN_OR_RAISE(auto spec, Spec(spec_id));
+      ICEBERG_ASSIGN_OR_RAISE(auto written_manifests,
+                              WriteDataManifests(data_files, spec));
+      new_manifests_.insert(new_manifests_.end(), written_manifests.begin(),
+                            written_manifests.end());
+    }
+    has_new_files_ = false;
+  }
+
+  return new_manifests_;
+}
+
+void FastAppend::AddFileToSummary(const std::shared_ptr<PartitionSpec>& spec,
+                                  const std::shared_ptr<DataFile>& file) {
+  // Update added data files count
+  auto it = summary_.find(SnapshotSummaryFields::kAddedDataFiles);
+  int64_t added_data_files = (it != summary_.end()) ? std::stoll(it->second) : 0;
+  summary_[SnapshotSummaryFields::kAddedDataFiles] = std::to_string(added_data_files + 1);
+
+  // Update added records count
+  it = summary_.find(SnapshotSummaryFields::kAddedRecords);
+  int64_t added_records = (it != summary_.end()) ? std::stoll(it->second) : 0;
+  summary_[SnapshotSummaryFields::kAddedRecords] =
+      std::to_string(added_records + file->record_count);
+
+  // Update added file size
+  it = summary_.find(SnapshotSummaryFields::kAddedFileSize);
+  int64_t added_file_size = (it != summary_.end()) ? std::stoll(it->second) : 0;
+  summary_[SnapshotSummaryFields::kAddedFileSize] =
+      std::to_string(added_file_size + file->file_size_in_bytes);
+
+  // Note: Partition summary is built during manifest writing, not here.
+  // The partition summary limit is handled in the manifest writer.
+}
+
+void FastAppend::AddManifestToSummary(const ManifestFile& manifest) {
+  // Update added data files count from manifest
+  if (manifest.added_files_count.has_value()) {
+    auto it = summary_.find(SnapshotSummaryFields::kAddedDataFiles);
+    int64_t added_data_files = (it != summary_.end()) ? std::stoll(it->second) : 0;
+    summary_[SnapshotSummaryFields::kAddedDataFiles] =
+        std::to_string(added_data_files + manifest.added_files_count.value());
+  }
+
+  // Update added records count from manifest
+  if (manifest.added_rows_count.has_value()) {
+    auto it = summary_.find(SnapshotSummaryFields::kAddedRecords);
+    int64_t added_records = (it != summary_.end()) ? std::stoll(it->second) : 0;
+    summary_[SnapshotSummaryFields::kAddedRecords] =
+        std::to_string(added_records + manifest.added_rows_count.value());
+  }
+}
+
+}  // namespace iceberg