src/iceberg/test/delete_loader_test.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/data/delete_loader.h"

#include <optional>
#include <utility>

#include <arrow/array.h>
#include <arrow/c/bridge.h>
#include <arrow/json/from_string.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include "iceberg/arrow/arrow_io_internal.h"
#include "iceberg/data/equality_delete_writer.h"
#include "iceberg/data/position_delete_writer.h"
#include "iceberg/deletes/dv_writer.h"
#include "iceberg/deletes/position_delete_index.h"
#include "iceberg/file_format.h"
#include "iceberg/file_io.h"
#include "iceberg/manifest/manifest_entry.h"
#include "iceberg/parquet/parquet_register.h"
#include "iceberg/partition_spec.h"
#include "iceberg/result.h"
#include "iceberg/row/partition_values.h"
#include "iceberg/schema.h"
#include "iceberg/schema_field.h"
#include "iceberg/schema_internal.h"
#include "iceberg/test/matchers.h"
#include "iceberg/type.h"
#include "iceberg/util/macros.h"
#include "iceberg/util/struct_like_set.h"

namespace iceberg {

class DeleteLoaderTest : public ::testing::Test {
 protected:
  static void SetUpTestSuite() { parquet::RegisterAll(); }

  void SetUp() override {
    file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO();
    schema_ = std::make_shared<Schema>(
        std::vector<SchemaField>{SchemaField::MakeRequired(1, "id", int32()),
                                 SchemaField::MakeOptional(2, "name", string())});
    partition_spec_ = PartitionSpec::Unpartitioned();
    loader_ = std::make_unique<DeleteLoader>(file_io_);
  }

  /// Write position deletes using WriteDelete API and return the DataFile metadata.
  std::shared_ptr<DataFile> WritePositionDeletes(
      const std::string& path,
      const std::vector<std::pair<std::string, int64_t>>& deletes) {
    PositionDeleteWriterOptions options{
        .path = path,
        .schema = schema_,
        .spec = partition_spec_,
        .partition = PartitionValues{},
        .format = FileFormatType::kParquet,
        .io = file_io_,
        .flush_threshold = 10000,
        .properties = {{"write.parquet.compression-codec", "uncompressed"}},
    };

    auto writer = PositionDeleteWriter::Make(options).value();
    for (const auto& [file_path, pos] : deletes) {
      ICEBERG_THROW_NOT_OK(writer->WriteDelete(file_path, pos));
    }
    ICEBERG_THROW_NOT_OK(writer->Close());
    return writer->Metadata().value().data_files[0];
  }

  /// Write equality deletes and return the DataFile metadata.
  std::shared_ptr<DataFile> WriteEqualityDeletes(
      const std::string& path, const std::string& json_data,
      std::vector<int32_t> equality_field_ids = {1, 2}) {
    EqualityDeleteWriterOptions options{
        .path = path,
        .schema = schema_,
        .spec = partition_spec_,
        .partition = PartitionValues{},
        .format = FileFormatType::kParquet,
        .io = file_io_,
        .equality_field_ids = std::move(equality_field_ids),
        .properties = {{"write.parquet.compression-codec", "uncompressed"}},
    };

    auto writer = EqualityDeleteWriter::Make(options).value();

    ArrowSchema arrow_c_schema;
    ICEBERG_THROW_NOT_OK(ToArrowSchema(*schema_, &arrow_c_schema));
    auto arrow_type = ::arrow::ImportType(&arrow_c_schema).ValueOrDie();
    auto test_data = ::arrow::json::ArrayFromJSONString(
                         ::arrow::struct_(arrow_type->fields()), json_data)
                         .ValueOrDie();
    ArrowArray arrow_array;
    auto export_status = ::arrow::ExportArray(*test_data, &arrow_array);
    if (!export_status.ok()) throw std::runtime_error(export_status.ToString());
    ICEBERG_THROW_NOT_OK(writer->Write(&arrow_array));
    ICEBERG_THROW_NOT_OK(writer->Close());
    return writer->Metadata().value().data_files[0];
  }

  /// Write a deletion vector to a Puffin file and return the DataFile metadata.
  std::shared_ptr<DataFile> WriteDeletionVector(const std::string& path,
                                                const std::string& referenced_data_file,
                                                const std::vector<int64_t>& positions) {
    auto writer =
        DVWriter::Make(DVWriterOptions{.path = path,
                                       .io = file_io_,
                                       .load_previous_deletes = [](std::string_view)
                                           -> Result<std::optional<PositionDeleteIndex>> {
                                         return std::nullopt;
                                       }})
            .value();
    for (int64_t pos : positions) {
      ICEBERG_THROW_NOT_OK(
          writer->Delete(referenced_data_file, pos, partition_spec_, PartitionValues{}));
    }
    ICEBERG_THROW_NOT_OK(writer->Close());
    return writer->Metadata().value().data_files[0];
  }

  std::shared_ptr<FileIO> file_io_;
  std::shared_ptr<Schema> schema_;
  std::shared_ptr<PartitionSpec> partition_spec_;
  std::unique_ptr<DeleteLoader> loader_;
};

// --- Position Delete Tests ---

TEST_F(DeleteLoaderTest, LoadPositionDeletesEmpty) {
  std::vector<std::shared_ptr<DataFile>> empty;
  auto result = loader_->LoadPositionDeletes(empty, "data.parquet");
  ASSERT_THAT(result, IsOk());
  ASSERT_TRUE(result->IsEmpty());
}

TEST_F(DeleteLoaderTest, LoadPositionDeletesSingleFile) {
  auto delete_file = WritePositionDeletes(
      "pos_deletes_1.parquet",
      {{"data.parquet", 0}, {"data.parquet", 5}, {"data.parquet", 10}});

  std::vector<std::shared_ptr<DataFile>> files = {delete_file};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsOk());

  auto& index = result.value();
  ASSERT_FALSE(index.IsEmpty());
  ASSERT_EQ(index.Cardinality(), 3);
  ASSERT_TRUE(index.IsDeleted(0));
  ASSERT_TRUE(index.IsDeleted(5));
  ASSERT_TRUE(index.IsDeleted(10));
  ASSERT_FALSE(index.IsDeleted(1));
  ASSERT_FALSE(index.IsDeleted(6));
}

TEST_F(DeleteLoaderTest, LoadPositionDeletesMultipleFiles) {
  auto file1 = WritePositionDeletes("pos_deletes_a.parquet",
                                    {{"data.parquet", 1}, {"data.parquet", 3}});
  auto file2 = WritePositionDeletes("pos_deletes_b.parquet",
                                    {{"data.parquet", 5}, {"data.parquet", 7}});

  std::vector<std::shared_ptr<DataFile>> files = {file1, file2};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsOk());

  auto& index = result.value();
  ASSERT_EQ(index.Cardinality(), 4);
  ASSERT_TRUE(index.IsDeleted(1));
  ASSERT_TRUE(index.IsDeleted(3));
  ASSERT_TRUE(index.IsDeleted(5));
  ASSERT_TRUE(index.IsDeleted(7));
}

TEST_F(DeleteLoaderTest, LoadPositionDeletesFiltersByDataFilePath) {
  auto delete_file =
      WritePositionDeletes("pos_deletes_mixed.parquet", {{"data_a.parquet", 0},
                                                         {"data_a.parquet", 2},
                                                         {"data_b.parquet", 10},
                                                         {"data_b.parquet", 20}});

  // Mixed paths keep the loader on the per-row filter path.
  ASSERT_FALSE(delete_file->referenced_data_file.has_value());

  std::vector<std::shared_ptr<DataFile>> files = {delete_file};

  // Load only positions for data_a.parquet
  auto result_a = loader_->LoadPositionDeletes(files, "data_a.parquet");
  ASSERT_THAT(result_a, IsOk());
  ASSERT_EQ(result_a.value().Cardinality(), 2);
  ASSERT_TRUE(result_a.value().IsDeleted(0));
  ASSERT_TRUE(result_a.value().IsDeleted(2));
  ASSERT_FALSE(result_a.value().IsDeleted(10));

  // Load only positions for data_b.parquet
  auto result_b = loader_->LoadPositionDeletes(files, "data_b.parquet");
  ASSERT_THAT(result_b, IsOk());
  ASSERT_EQ(result_b.value().Cardinality(), 2);
  ASSERT_TRUE(result_b.value().IsDeleted(10));
  ASSERT_TRUE(result_b.value().IsDeleted(20));
  ASSERT_FALSE(result_b.value().IsDeleted(0));

  // Load for non-existent file
  auto result_none = loader_->LoadPositionDeletes(files, "nonexistent.parquet");
  ASSERT_THAT(result_none, IsOk());
  ASSERT_TRUE(result_none.value().IsEmpty());
}

// The loaded index keeps the source Parquet position delete file, so a DV
// writer can report file-scoped deletes as rewritten.
TEST_F(DeleteLoaderTest, LoadPositionDeletesTracksSourceDeleteFile) {
  // Single referenced data file -> writer sets the hint -> fast path.
  auto file_scoped = WritePositionDeletes("pos_src_scoped.parquet",
                                          {{"data.parquet", 0}, {"data.parquet", 5}});
  ASSERT_TRUE(file_scoped->referenced_data_file.has_value());

  std::vector<std::shared_ptr<DataFile>> files = {file_scoped};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsOk());
  ASSERT_EQ(result.value().delete_files().size(), 1u);
  EXPECT_EQ(result.value().delete_files()[0], file_scoped);
}

TEST_F(DeleteLoaderTest, LoadPositionDeletesTracksSourceForFilterPath) {
  auto mixed = WritePositionDeletes("pos_src_mixed.parquet",
                                    {{"data_a.parquet", 0}, {"data_b.parquet", 10}});
  ASSERT_FALSE(mixed->referenced_data_file.has_value());

  std::vector<std::shared_ptr<DataFile>> files = {mixed};

  auto matched = loader_->LoadPositionDeletes(files, "data_a.parquet");
  ASSERT_THAT(matched, IsOk());
  ASSERT_EQ(matched.value().delete_files().size(), 1u);
  EXPECT_EQ(matched.value().delete_files()[0], mixed);

  auto unmatched = loader_->LoadPositionDeletes(files, "data_c.parquet");
  ASSERT_THAT(unmatched, IsOk());
  ASSERT_TRUE(unmatched.value().IsEmpty());
  ASSERT_EQ(unmatched.value().delete_files().size(), 1u);
  EXPECT_EQ(unmatched.value().delete_files()[0], mixed);
}

TEST_F(DeleteLoaderTest, LoadPositionDeletesRejectsNullFile) {
  std::vector<std::shared_ptr<DataFile>> files = {nullptr};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
}

TEST_F(DeleteLoaderTest, LoadPositionDeletesSkipsMismatchedReferencedDataFile) {
  auto delete_file = std::make_shared<DataFile>(DataFile{
      .content = DataFile::Content::kPositionDeletes,
      .file_path = "missing-pos-delete.parquet",
      .file_format = FileFormatType::kParquet,
      .referenced_data_file = "other-data.parquet",
  });

  std::vector<std::shared_ptr<DataFile>> files = {delete_file};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsOk());
  ASSERT_TRUE(result.value().IsEmpty());
}

TEST_F(DeleteLoaderTest, LoadPositionDeletesFastPathHonorsReferencedDataFile) {
  // Large enough to exercise the bulk path.
  constexpr int64_t kRowCount = 128;
  std::vector<std::pair<std::string, int64_t>> deletes;
  deletes.reserve(kRowCount);
  for (int64_t i = 0; i < kRowCount; ++i) {
    deletes.emplace_back("data.parquet", i);
  }
  auto delete_file = WritePositionDeletes("pos_deletes_fast_path.parquet", deletes);

  ASSERT_TRUE(delete_file->referenced_data_file.has_value());
  ASSERT_EQ(delete_file->referenced_data_file.value(), "data.parquet");

  std::vector<std::shared_ptr<DataFile>> files = {delete_file};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsOk());

  auto& index = result.value();
  ASSERT_EQ(index.Cardinality(), kRowCount);
  ASSERT_TRUE(index.IsDeleted(0));
  ASSERT_TRUE(index.IsDeleted(kRowCount / 2));
  ASSERT_TRUE(index.IsDeleted(kRowCount - 1));
  ASSERT_FALSE(index.IsDeleted(kRowCount));
}

TEST_F(DeleteLoaderTest, LoadDeletionVector) {
  auto dv_file =
      WriteDeletionVector("dv-a.puffin", "data.parquet", {0, 5, 10, 4'000'000'000LL});

  std::vector<std::shared_ptr<DataFile>> files = {dv_file};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsOk());

  auto& index = result.value();
  ASSERT_EQ(index.Cardinality(), 4);
  ASSERT_TRUE(index.IsDeleted(0));
  ASSERT_TRUE(index.IsDeleted(5));
  ASSERT_TRUE(index.IsDeleted(10));
  ASSERT_TRUE(index.IsDeleted(4'000'000'000LL));
  ASSERT_FALSE(index.IsDeleted(1));
}

TEST_F(DeleteLoaderTest, LoadDeletionVectorRejectsMismatchedReferencedDataFile) {
  // A DV referencing another data file must not be silently skipped: the single
  // DV applies to its referenced file, so a mismatch is bad grouping/metadata.
  auto dv_file = WriteDeletionVector("dv-b.puffin", "other-data.parquet", {1, 2, 3});

  std::vector<std::shared_ptr<DataFile>> files = {dv_file};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
}

TEST_F(DeleteLoaderTest, LoadDeletionVectorRequiresContentOffsetAndSize) {
  auto dv_file = std::make_shared<DataFile>(DataFile{
      .content = DataFile::Content::kPositionDeletes,
      .file_path = "dv.puffin",
      .file_format = FileFormatType::kPuffin,
      .referenced_data_file = "data.parquet",
  });
  std::vector<std::shared_ptr<DataFile>> files = {dv_file};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
}

TEST_F(DeleteLoaderTest, LoadDeletionVectorRejectsCardinalityMismatch) {
  auto dv_file = WriteDeletionVector("dv-card.puffin", "data.parquet", {0, 1, 2});
  // Corrupt the recorded cardinality so it no longer matches the bitmap.
  dv_file->record_count = 99;

  std::vector<std::shared_ptr<DataFile>> files = {dv_file};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
}

// End to end: a previously written Parquet position delete is loaded, merged
// into a new DV, and reported in rewritten_delete_files when it is file-scoped.
TEST_F(DeleteLoaderTest, DVWriterReportsFileScopedPositionDeleteAsRewritten) {
  // A file-scoped position delete (single referenced data file).
  auto previous = WritePositionDeletes("prev_pos.parquet",
                                       {{"data.parquet", 100}, {"data.parquet", 200}});
  ASSERT_TRUE(previous->referenced_data_file.has_value());

  auto* loader = loader_.get();
  ICEBERG_UNWRAP_OR_FAIL(
      auto writer,
      DVWriter::Make(DVWriterOptions{
          .path = "dv-replace.puffin",
          .io = file_io_,
          .load_previous_deletes =
              [&](std::string_view path) -> Result<std::optional<PositionDeleteIndex>> {
            std::vector<std::shared_ptr<DataFile>> files = {previous};
            ICEBERG_ASSIGN_OR_RAISE(auto index, loader->LoadPositionDeletes(files, path));
            if (index.IsEmpty()) {
              return std::nullopt;
            }
            return std::optional<PositionDeleteIndex>(std::move(index));
          },
      }));

  ICEBERG_THROW_NOT_OK(
      writer->Delete("data.parquet", 0, partition_spec_, PartitionValues{}));
  ICEBERG_THROW_NOT_OK(writer->Close());

  ICEBERG_UNWRAP_OR_FAIL(auto result, writer->Metadata());
  ASSERT_EQ(result.data_files.size(), 1u);
  // New position plus the two previous positions.
  EXPECT_EQ(result.data_files[0]->record_count, 3);
  // The previous Parquet position delete is file-scoped and must be reported.
  ASSERT_EQ(result.rewritten_delete_files.size(), 1u);
  EXPECT_EQ(result.rewritten_delete_files[0], previous);
}

TEST_F(DeleteLoaderTest, LoadPositionDeletesRejectsWrongContent) {
  auto data_file = std::make_shared<DataFile>(DataFile{
      .content = DataFile::Content::kData,
      .file_path = "data.parquet",
      .file_format = FileFormatType::kParquet,
  });
  std::vector<std::shared_ptr<DataFile>> files = {data_file};
  auto result = loader_->LoadPositionDeletes(files, "data.parquet");
  ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
}

// --- Equality Delete Tests ---

TEST_F(DeleteLoaderTest, LoadEqualityDeletesEmpty) {
  std::vector<std::shared_ptr<DataFile>> empty;
  StructType eq_type(std::vector<SchemaField>{
      SchemaField::MakeRequired(1, "id", int32()),
      SchemaField::MakeOptional(2, "name", string()),
  });

  auto result = loader_->LoadEqualityDeletes(empty, eq_type);
  ASSERT_THAT(result, IsOk());
  ASSERT_TRUE(result.value()->IsEmpty());
  ASSERT_EQ(result.value()->Size(), 0);
}

TEST_F(DeleteLoaderTest, LoadEqualityDeletesSingleFile) {
  auto delete_file =
      WriteEqualityDeletes("eq_deletes_1.parquet", R"([[1, "Alice"], [2, "Bob"]])");

  StructType eq_type(std::vector<SchemaField>{
      SchemaField::MakeRequired(1, "id", int32()),
      SchemaField::MakeOptional(2, "name", string()),
  });

  std::vector<std::shared_ptr<DataFile>> files = {delete_file};
  auto result = loader_->LoadEqualityDeletes(files, eq_type);
  ASSERT_THAT(result, IsOk());

  ASSERT_FALSE(result.value()->IsEmpty());
  ASSERT_EQ(result.value()->Size(), 2);
}

TEST_F(DeleteLoaderTest, LoadEqualityDeletesMultipleFiles) {
  auto file1 = WriteEqualityDeletes("eq_deletes_a.parquet", R"([[1, "Alice"]])");
  auto file2 =
      WriteEqualityDeletes("eq_deletes_b.parquet", R"([[2, "Bob"], [3, "Charlie"]])");

  StructType eq_type(std::vector<SchemaField>{
      SchemaField::MakeRequired(1, "id", int32()),
      SchemaField::MakeOptional(2, "name", string()),
  });

  std::vector<std::shared_ptr<DataFile>> files = {file1, file2};
  auto result = loader_->LoadEqualityDeletes(files, eq_type);
  ASSERT_THAT(result, IsOk());
  ASSERT_EQ(result.value()->Size(), 3);
}

TEST_F(DeleteLoaderTest, LoadEqualityDeletesRejectsWrongContent) {
  auto data_file = std::make_shared<DataFile>(DataFile{
      .content = DataFile::Content::kData,
      .file_path = "data.parquet",
      .file_format = FileFormatType::kParquet,
  });
  StructType eq_type(std::vector<SchemaField>{
      SchemaField::MakeRequired(1, "id", int32()),
  });

  std::vector<std::shared_ptr<DataFile>> files = {data_file};
  auto result = loader_->LoadEqualityDeletes(files, eq_type);
  ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
}

}  // namespace iceberg