src/iceberg/update/snapshot_update.cc

/*
 * 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/snapshot_update.h"

#include <algorithm>
#include <format>
#include <ranges>
#include <span>

#include "iceberg/constants.h"
#include "iceberg/file_io.h"
#include "iceberg/inheritable_metadata.h"
#include "iceberg/logging/log_macros.h"
#include "iceberg/manifest/manifest_entry.h"
#include "iceberg/manifest/manifest_list.h"
#include "iceberg/manifest/manifest_reader.h"
#include "iceberg/manifest/manifest_writer.h"
#include "iceberg/manifest/rolling_manifest_writer.h"
#include "iceberg/metrics/commit_report.h"
#include "iceberg/metrics/metrics_context.h"
#include "iceberg/metrics/metrics_reporter.h"
#include "iceberg/partition_summary_internal.h"
#include "iceberg/table.h"  // IWYU pragma: keep
#include "iceberg/transaction.h"
#include "iceberg/util/error_util_internal.h"
#include "iceberg/util/executor_util_internal.h"
#include "iceberg/util/macros.h"
#include "iceberg/util/snapshot_util.h"
#include "iceberg/util/string_util.h"
#include "iceberg/util/task_group.h"
#include "iceberg/util/uuid.h"

namespace iceberg {

namespace {

Status UpdateTotal(std::unordered_map<std::string, std::string>& summary,
                   const std::unordered_map<std::string, std::string>& previous_summary,
                   const std::string& total_property, const std::string& added_property,
                   const std::string& deleted_property) {
  auto total_it = previous_summary.find(total_property);
  if (total_it != previous_summary.end()) {
    auto parsed_total = StringUtils::ParseNumber<int64_t>(total_it->second);
    if (!parsed_total.has_value()) {
      return {};
    }
    int64_t new_total = parsed_total.value();

    auto added_it = summary.find(added_property);
    if (new_total >= 0 && added_it != summary.end()) {
      auto parsed_added = StringUtils::ParseNumber<int64_t>(added_it->second);
      if (!parsed_added.has_value()) {
        return {};
      }
      new_total += parsed_added.value();
    }

    auto deleted_it = summary.find(deleted_property);
    if (new_total >= 0 && deleted_it != summary.end()) {
      auto parsed_deleted = StringUtils::ParseNumber<int64_t>(deleted_it->second);
      if (!parsed_deleted.has_value()) {
        return {};
      }
      new_total -= parsed_deleted.value();
    }

    if (new_total >= 0) {
      summary[total_property] = std::to_string(new_total);
    }
  }
  return {};
}

constexpr size_t kMinManifestWriterGroupSize = 10'000;

template <typename T, typename WriteGroup>
Result<std::vector<ManifestFile>> WriteManifestGroups(OptionalExecutor executor,
                                                      int32_t max_parallelism,
                                                      std::span<const T> files,
                                                      WriteGroup&& write_group) {
  const auto limit = static_cast<int32_t>(
      (files.size() + kMinManifestWriterGroupSize / 2) / kMinManifestWriterGroupSize);
  const auto group_count =
      static_cast<size_t>(std::max<int32_t>(1, std::min(max_parallelism, limit)));
  const size_t group_size = (files.size() + group_count - 1) / group_count;
  // TODO(zehua): Replace the manual offset calculation with `std::views::chunk`
  // once the supported libc++ provides it.
  auto groups =
      std::views::iota(0UZ, group_count) |
      std::views::transform([files, group_size](size_t group_index) {
        const size_t offset = group_index * group_size;
        return files.subspan(offset, std::min(group_size, files.size() - offset));
      });
  return ParallelCollect(executor, groups, std::forward<WriteGroup>(write_group));
}

// Add metadata to a manifest file by reading it and extracting statistics.
Result<ManifestFile> AddMetadata(const ManifestFile& manifest, std::shared_ptr<FileIO> io,
                                 const TableMetadata& metadata) {
  ICEBERG_PRECHECK(!manifest.added_snapshot_id.has_value(),
                   "Manifest {} already has assigned a snapshot id: {}",
                   manifest.manifest_path, manifest.added_snapshot_id.value());

  ICEBERG_ASSIGN_OR_RAISE(auto schema, metadata.Schema());
  ICEBERG_ASSIGN_OR_RAISE(auto spec,
                          metadata.PartitionSpecById(manifest.partition_spec_id));
  ICEBERG_ASSIGN_OR_RAISE(auto partition_type, spec->PartitionType(*schema));

  ICEBERG_ASSIGN_OR_RAISE(auto inheritable_metadata, InheritableMetadataFactory::Empty());
  ICEBERG_ASSIGN_OR_RAISE(
      auto reader,
      ManifestReader::Make(manifest.manifest_path, manifest.manifest_length,
                           std::move(io), schema, spec, std::move(inheritable_metadata),
                           manifest.first_row_id));
  ICEBERG_ASSIGN_OR_RAISE(auto entries, reader->Entries());

  PartitionSummary stats(*partition_type);
  int32_t added_files = 0;
  int64_t added_rows = 0;
  int32_t existing_files = 0;
  int64_t existing_rows = 0;
  int32_t deleted_files = 0;
  int64_t deleted_rows = 0;

  std::optional<int64_t> snapshot_id;
  int64_t max_snapshot_id = std::numeric_limits<int64_t>::min();
  for (const auto& entry : entries) {
    ICEBERG_PRECHECK(entry.data_file != nullptr,
                     "Manifest entry in {} is missing data_file", manifest.manifest_path);

    if (entry.snapshot_id.has_value() && entry.snapshot_id.value() > max_snapshot_id) {
      max_snapshot_id = entry.snapshot_id.value();
    }

    switch (entry.status) {
      case ManifestStatus::kAdded: {
        added_files += 1;
        added_rows += entry.data_file->record_count;
        if (!snapshot_id.has_value() && entry.snapshot_id.has_value()) {
          snapshot_id = entry.snapshot_id;
        }
      } break;
      case ManifestStatus::kExisting: {
        existing_files += 1;
        existing_rows += entry.data_file->record_count;
      } break;
      case ManifestStatus::kDeleted: {
        deleted_files += 1;
        deleted_rows += entry.data_file->record_count;
        if (!snapshot_id.has_value() && entry.snapshot_id.has_value()) {
          snapshot_id = entry.snapshot_id;
        }
      } break;
    }

    ICEBERG_RETURN_UNEXPECTED(stats.Update(entry.data_file->partition));
  }

  if (!snapshot_id.has_value()) {
    // If no files were added or deleted, use the largest snapshot ID in the manifest
    snapshot_id = max_snapshot_id;
  }

  ICEBERG_ASSIGN_OR_RAISE(auto partition_summaries, stats.Summaries());

  ManifestFile enriched = manifest;
  enriched.added_snapshot_id = snapshot_id.value();
  enriched.added_files_count = added_files;
  enriched.existing_files_count = existing_files;
  enriched.deleted_files_count = deleted_files;
  enriched.added_rows_count = added_rows;
  enriched.existing_rows_count = existing_rows;
  enriched.deleted_rows_count = deleted_rows;
  enriched.partitions = std::move(partition_summaries);
  enriched.first_row_id = std::nullopt;
  return enriched;
}

}  // anonymous namespace

SnapshotUpdate::~SnapshotUpdate() = default;

SnapshotUpdate::SnapshotUpdate(std::shared_ptr<TransactionContext> ctx)
    : PendingUpdate(std::move(ctx)),
      can_inherit_snapshot_id_(
          base().format_version > 1 ||
          base().properties.Get(TableProperties::kSnapshotIdInheritanceEnabled)),
      commit_uuid_(Uuid::GenerateV7().ToString()),
      target_manifest_size_bytes_(
          base().properties.Get(TableProperties::kManifestTargetSizeBytes)),
      commit_metrics_(CommitMetrics::Make(*MetricsContext::Default())),
      reporter_(ctx_->table->reporter()) {}

Status SnapshotUpdate::Commit() {
  ICEBERG_RETURN_UNEXPECTED(CheckCommitAllowed());
  [[maybe_unused]] auto commit_timer = commit_metrics_->total_duration->Start();
  return PendingUpdate::Commit();
}

Status SnapshotUpdate::ReportCommit() const {
  ICEBERG_DCHECK(staged_snapshot_ != nullptr,
                 "Staged snapshot is null after a successful commit");

  if (!reporter_) {
    return {};
  }

  const auto operation = staged_snapshot_->Operation();
  CommitReport report{
      .table_name = ctx_->table->full_name(),
      .snapshot_id = staged_snapshot_->snapshot_id,
      .sequence_number = staged_snapshot_->sequence_number,
      .operation = operation.has_value() ? std::string(operation.value()) : "",
      .commit_metrics =
          CommitMetricsResult::From(*commit_metrics_, staged_snapshot_->summary),
      .metadata = {},
  };
  return reporter_->Report(report);
}

void SnapshotUpdate::SetSummaryProperty(const std::string& property,
                                        const std::string& value) {
  summary_.Set(property, value);
}

Result<std::vector<ManifestFile>> SnapshotUpdate::WriteDataManifests(
    std::span<const std::shared_ptr<DataFile>> files,
    const std::shared_ptr<PartitionSpec>& spec,
    std::optional<int64_t> data_sequence_number) {
  if (files.empty()) {
    return std::vector<ManifestFile>{};
  }

  ICEBERG_ASSIGN_OR_RAISE(auto current_schema, base().Schema());
  const int8_t format_version = base().format_version;
  const int64_t snapshot_id = SnapshotId();
  auto make_writer = [&]() {
    return ManifestWriter::MakeWriter(format_version, snapshot_id, ManifestPath(),
                                      ctx_->table->io(), spec, current_schema,
                                      ManifestContent::kData);
  };

  return WriteManifestGroups(
      write_manifest_executor_, write_manifest_parallelism_, files,
      [&](std::span<const std::shared_ptr<DataFile>> group)
          -> Result<std::vector<ManifestFile>> {
        RollingManifestWriter rolling_writer(make_writer, target_manifest_size_bytes_);
        for (const auto& file : group) {
          ICEBERG_RETURN_UNEXPECTED(
              rolling_writer.WriteAddedEntry(file, data_sequence_number));
        }
        ICEBERG_RETURN_UNEXPECTED(rolling_writer.Close());
        return rolling_writer.ToManifestFiles();
      });
}

Result<std::vector<ManifestFile>> SnapshotUpdate::WriteDeleteManifests(
    std::span<const ContentFileWithSequenceNumber> files,
    const std::shared_ptr<PartitionSpec>& spec) {
  if (files.empty()) {
    return std::vector<ManifestFile>{};
  }

  ICEBERG_ASSIGN_OR_RAISE(auto current_schema, base().Schema());
  const int8_t format_version = base().format_version;
  const int64_t snapshot_id = SnapshotId();
  auto make_writer = [&]() {
    return ManifestWriter::MakeWriter(format_version, snapshot_id, ManifestPath(),
                                      ctx_->table->io(), spec, current_schema,
                                      ManifestContent::kDeletes);
  };

  return WriteManifestGroups(
      write_manifest_executor_, write_manifest_parallelism_, files,
      [&](std::span<const ContentFileWithSequenceNumber> group)
          -> Result<std::vector<ManifestFile>> {
        RollingManifestWriter rolling_writer(make_writer, target_manifest_size_bytes_);
        for (const auto& entry : group) {
          ICEBERG_RETURN_UNEXPECTED(
              rolling_writer.WriteAddedEntry(entry.file, entry.data_sequence_number));
        }
        ICEBERG_RETURN_UNEXPECTED(rolling_writer.Close());
        return rolling_writer.ToManifestFiles();
      });
}

int64_t SnapshotUpdate::SnapshotId() {
  if (!snapshot_id_.has_value()) {
    snapshot_id_ = SnapshotUtil::GenerateSnapshotId(base());
  }
  return snapshot_id_.value();
}

Result<SnapshotUpdate::ApplyResult> SnapshotUpdate::Apply() {
  commit_metrics_->attempts->Increment();
  ICEBERG_RETURN_UNEXPECTED(CheckErrors());

  ICEBERG_ASSIGN_OR_RAISE(auto parent_snapshot,
                          SnapshotUtil::OptionalLatestSnapshot(base(), target_branch_));

  int64_t sequence_number = base().NextSequenceNumber();
  std::optional<int64_t> parent_snapshot_id =
      parent_snapshot ? std::make_optional(parent_snapshot->snapshot_id) : std::nullopt;

  ICEBERG_RETURN_UNEXPECTED(Validate(base(), parent_snapshot));

  ICEBERG_ASSIGN_OR_RAISE(auto manifests, Apply(base(), parent_snapshot));
  auto metadata_tasks = TaskGroup().SetExecutor(plan_executor_);
  for (auto& manifest : manifests) {
    if (manifest.added_snapshot_id.has_value()) {
      continue;
    }
    metadata_tasks.Submit([&manifest, this]() -> Status {
      ICEBERG_ASSIGN_OR_RAISE(manifest, AddMetadata(manifest, ctx_->table->io(), base()));
      return {};
    });
  }
  ICEBERG_RETURN_UNEXPECTED(std::move(metadata_tasks).Run());

  std::string manifest_list_path = ManifestListPath();
  ICEBERG_ASSIGN_OR_RAISE(
      auto writer, ManifestListWriter::MakeWriter(base().format_version, SnapshotId(),
                                                  parent_snapshot_id, manifest_list_path,
                                                  ctx_->table->io(), sequence_number,
                                                  base().next_row_id));
  ICEBERG_RETURN_UNEXPECTED(writer->AddAll(manifests));
  ICEBERG_RETURN_UNEXPECTED(writer->Close());

  std::optional<int64_t> next_row_id;
  std::optional<int64_t> assigned_rows;
  if (base().format_version >= 3) {
    ICEBERG_CHECK(writer->next_row_id().has_value(),
                  "row id is required by format version >= 3");
    next_row_id = base().next_row_id;
    assigned_rows = writer->next_row_id().value() - base().next_row_id;
  }

  std::string op = operation();
  ICEBERG_CHECK(!op.empty(), "Snapshot operation cannot be empty");

  if (op == DataOperation::kReplace) {
    const auto summary = Summary();
    auto added_records_it = summary.find(SnapshotSummaryFields::kAddedRecords);
    auto replaced_records_it = summary.find(SnapshotSummaryFields::kDeletedRecords);
    if (added_records_it != summary.cend() && replaced_records_it != summary.cend()) {
      ICEBERG_ASSIGN_OR_RAISE(auto added_records, StringUtils::ParseNumber<int64_t>(
                                                      added_records_it->second));
      ICEBERG_ASSIGN_OR_RAISE(auto replaced_records, StringUtils::ParseNumber<int64_t>(
                                                         replaced_records_it->second));
      ICEBERG_PRECHECK(
          added_records <= replaced_records,
          "Invalid REPLACE operation: {} added records > {} replaced records",
          added_records, replaced_records);
    }
  }

  ICEBERG_ASSIGN_OR_RAISE(auto summary, ComputeSummary(base()));
  ICEBERG_ASSIGN_OR_RAISE(
      staged_snapshot_,
      Snapshot::Make(sequence_number, SnapshotId(), parent_snapshot_id,
                     CurrentTimePointMs(), std::move(op), std::move(summary),
                     base().current_schema_id, std::move(manifest_list_path), next_row_id,
                     assigned_rows));

  return ApplyResult{.snapshot = staged_snapshot_,
                     .target_branch = target_branch_,
                     .stage_only = stage_only_};
}

Status SnapshotUpdate::Finalize([[maybe_unused]] const TableMetadata& metadata) {
  if (staged_snapshot_ == nullptr) {
    return {};
  }

  // Only files created by this update are tracked, and committed paths are kept below.
  internal::LogAndIgnoreFailure("Snapshot cleanup", [this]() -> Status {
    auto snapshot_reader = SnapshotReader(staged_snapshot_.get());
    ICEBERG_ASSIGN_OR_RAISE(auto manifests, snapshot_reader.Manifests(ctx_->table->io()));
    auto committed = manifests | std::views::transform([](const auto& manifest) {
                       return manifest.manifest_path;
                     }) |
                     std::ranges::to<std::unordered_set<std::string>>();
    internal::LogAndIgnoreFailure("Snapshot operation cleanup", [this, &committed] {
      return CleanUncommitted(committed);
    });
    committed.insert(staged_snapshot_->manifest_list);
    std::vector<std::string> unused;
    {
      std::lock_guard lock(staging_mutex_);
      for (const auto& path : staged_files_) {
        if (!committed.contains(path)) {
          unused.push_back(path);
        }
      }
    }
    for (const auto& path : unused) {
      std::ignore = DeleteFile(path);
    }
    return {};
  });

  {
    std::lock_guard lock(staging_mutex_);
    staged_files_.clear();
  }
  auto report_status = ReportCommit();
  staged_snapshot_.reset();
  summary_.Clear();
  return report_status;
}

Result<std::unordered_map<std::string, std::string>> SnapshotUpdate::ComputeSummary(
    const TableMetadata& previous) {
  std::unordered_map<std::string, std::string> summary = Summary();
  if (summary.empty()) {
    return summary;
  }

  // Get previous summary from the target branch
  std::unordered_map<std::string, std::string> previous_summary;
  if (auto ref_it = previous.refs.find(target_branch_); ref_it != previous.refs.end()) {
    if (auto snap_it = previous.SnapshotById(ref_it->second->snapshot_id);
        snap_it.has_value()) {
      previous_summary = snap_it.value()->summary;
    }
  } else {
    // if there was no previous snapshot, default the summary to start totals at 0
    previous_summary[SnapshotSummaryFields::kTotalRecords] = "0";
    previous_summary[SnapshotSummaryFields::kTotalFileSize] = "0";
    previous_summary[SnapshotSummaryFields::kTotalDataFiles] = "0";
    previous_summary[SnapshotSummaryFields::kTotalDeleteFiles] = "0";
    previous_summary[SnapshotSummaryFields::kTotalPosDeletes] = "0";
    previous_summary[SnapshotSummaryFields::kTotalEqDeletes] = "0";
  }

  // Update totals
  ICEBERG_RETURN_UNEXPECTED(UpdateTotal(
      summary, previous_summary, SnapshotSummaryFields::kTotalRecords,
      SnapshotSummaryFields::kAddedRecords, SnapshotSummaryFields::kDeletedRecords));
  ICEBERG_RETURN_UNEXPECTED(UpdateTotal(
      summary, previous_summary, SnapshotSummaryFields::kTotalFileSize,
      SnapshotSummaryFields::kAddedFileSize, SnapshotSummaryFields::kRemovedFileSize));
  ICEBERG_RETURN_UNEXPECTED(UpdateTotal(
      summary, previous_summary, SnapshotSummaryFields::kTotalDataFiles,
      SnapshotSummaryFields::kAddedDataFiles, SnapshotSummaryFields::kDeletedDataFiles));
  ICEBERG_RETURN_UNEXPECTED(UpdateTotal(summary, previous_summary,
                                        SnapshotSummaryFields::kTotalDeleteFiles,
                                        SnapshotSummaryFields::kAddedDeleteFiles,
                                        SnapshotSummaryFields::kRemovedDeleteFiles));
  ICEBERG_RETURN_UNEXPECTED(UpdateTotal(summary, previous_summary,
                                        SnapshotSummaryFields::kTotalPosDeletes,
                                        SnapshotSummaryFields::kAddedPosDeletes,
                                        SnapshotSummaryFields::kRemovedPosDeletes));
  ICEBERG_RETURN_UNEXPECTED(UpdateTotal(
      summary, previous_summary, SnapshotSummaryFields::kTotalEqDeletes,
      SnapshotSummaryFields::kAddedEqDeletes, SnapshotSummaryFields::kRemovedEqDeletes));

  // TODO(xxx): add custom summary fields like engine info
  return summary;
}

Status SnapshotUpdate::CleanStaged() {
  internal::LogAndIgnoreFailure("Snapshot staging cleanup",
                                [this] { return CleanUncommitted({}); });
  // Retry paths that operation-specific cleanup did not delete.
  std::vector<std::string> paths;
  {
    std::lock_guard lock(staging_mutex_);
    paths.assign(staged_files_.begin(), staged_files_.end());
  }
  for (const auto& path : paths) {
    std::ignore = DeleteFile(path);
  }
  staged_snapshot_.reset();
  summary_.Clear();
  return {};
}

void SnapshotUpdate::RegisterStagedFile(const std::string& path) {
  std::lock_guard lock(staging_mutex_);
  staged_files_.insert(path);
}

void SnapshotUpdate::UnregisterStagedFile(const std::string& path) {
  std::lock_guard lock(staging_mutex_);
  staged_files_.erase(path);
}

Status SnapshotUpdate::DeleteFile(const std::string& path) noexcept {
  try {
    auto result = delete_func_ ? delete_func_(path) : ctx_->table->io()->DeleteFile(path);
    if (!result && result.error().kind != ErrorKind::kNotFound) {
      RegisterStagedFile(path);
      ICEBERG_LOG_WARN("Cannot clean staged file {}: {}", path, result.error().message);
      return {};
    }
    std::lock_guard lock(staging_mutex_);
    staged_files_.erase(path);
  } catch (const std::exception& e) {
    RegisterStagedFile(path);
    ICEBERG_LOG_WARN("Cannot clean staged file {}: {}", path, e.what());
  } catch (...) {
    RegisterStagedFile(path);
    ICEBERG_LOG_WARN("Cannot clean staged file {}: unknown exception", path);
  }
  return {};
}

std::string SnapshotUpdate::ManifestListPath() {
  // Generate manifest list path
  // Format: {metadata_location}/snap-{snapshot_id}-{attempt}-{uuid}.avro
  int64_t snapshot_id = SnapshotId();
  auto attempt = attempt_.fetch_add(1, std::memory_order_relaxed) + 1;
  std::string filename =
      std::format("snap-{}-{}-{}.avro", snapshot_id, attempt, commit_uuid_);
  auto path = ctx_->MetadataFileLocation(filename);
  RegisterStagedFile(path);
  return path;
}

SnapshotSummaryBuilder SnapshotUpdate::BuildManifestCountSummary(
    std::span<const ManifestFile> manifests, int32_t replaced_manifests_count) {
  SnapshotSummaryBuilder summary;
  int32_t manifests_created = 0;
  int32_t manifests_kept = 0;
  int64_t snapshot_id = SnapshotId();
  for (const auto& manifest : manifests) {
    if (manifest.added_snapshot_id == snapshot_id) {
      ++manifests_created;
    } else if (manifest.added_snapshot_id.has_value()) {
      ++manifests_kept;
    }
  }

  summary.Set(SnapshotSummaryFields::kManifestsCreated,
              std::to_string(manifests_created));
  summary.Set(SnapshotSummaryFields::kManifestsKept, std::to_string(manifests_kept));
  summary.Set(SnapshotSummaryFields::kManifestsReplaced,
              std::to_string(replaced_manifests_count));
  return summary;
}

std::string SnapshotUpdate::ManifestPath() {
  // Generate manifest path
  // Format: {metadata_location}/{uuid}-m{manifest_count}.avro
  auto manifest_count = manifest_count_.fetch_add(1, std::memory_order_relaxed);
  std::string filename = std::format("{}-m{}.avro", commit_uuid_, manifest_count);
  auto path = ctx_->MetadataFileLocation(filename);
  RegisterStagedFile(path);
  return path;
}

}  // namespace iceberg