Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 92 additions & 9 deletions src/iceberg/data/position_delete_writer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,11 @@

#include "iceberg/data/position_delete_writer.h"

#include <functional>
#include <map>
#include <set>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

Expand Down Expand Up @@ -57,15 +60,68 @@ class PositionDeleteWriter::Impl {
ICEBERG_ASSIGN_OR_RAISE(auto writer,
WriterFactoryRegistry::Open(options.format, writer_options));

return std::unique_ptr<Impl>(
auto impl = std::unique_ptr<Impl>(
new Impl(std::move(options), std::move(delete_schema), std::move(writer)));
ICEBERG_RETURN_UNEXPECTED(impl->InitSchema());
return impl;
}

~Impl() {
ArrowArrayViewReset(&array_view_);
if (arrow_schema_.release != nullptr) {
ArrowSchemaRelease(&arrow_schema_);
}
}

Status Write(ArrowArray* data) {
ICEBERG_PRECHECK(data != nullptr, "Position delete data must not be null");
internal::ArrowArrayGuard data_guard(data);
ICEBERG_PRECHECK(data->offset == 0,
"Position delete data with a non-zero offset is not supported");
ICEBERG_PRECHECK(buffered_paths_.empty(),
"Cannot write batch data when there are buffered deletes.");
// TODO(anyone): Extract file paths from ArrowArray to update referenced_paths_.
return writer_->Write(data);

ArrowError error;
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
ArrowArrayViewSetArray(&array_view_, data, &error), error);

const auto* path_view = array_view_.children[0];
const auto* pos_view = array_view_.children[1];

// A batch usually references files that earlier batches already referenced, so
// record paths optimistically: a known path costs a lookup instead of a scratch
// allocation. Entries added by this batch are rolled back unless the write
// succeeds, so a rejected batch still leaves no trace in the metadata.
pending_references_.clear();
for (int64_t i = 0; i < data->length; ++i) {
if (ArrowArrayViewIsNull(path_view, i)) {
RollbackPendingReferences();
return InvalidArrowData(
"Position delete file paths must not contain null values");
}
if (ArrowArrayViewIsNull(pos_view, i)) {
RollbackPendingReferences();
return InvalidArrowData("Position delete positions must not contain null values");
}
auto path = ArrowArrayViewGetStringUnsafe(path_view, i);
if (path.size_bytes == 0) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we error out in this case?

RollbackPendingReferences();
return InvalidArrowData("Position delete file paths must not be empty");
}
std::string_view file_path(path.data, static_cast<size_t>(path.size_bytes));
if (!referenced_paths_.contains(file_path)) {
pending_references_.push_back(
referenced_paths_.insert(std::string(file_path)).first);
}
}

Status status = writer_->Write(data);
if (!status) {
RollbackPendingReferences();
return status;
}
pending_references_.clear();
return {};
}

Status WriteDelete(std::string_view file_path, int64_t pos) {
Expand Down Expand Up @@ -163,6 +219,8 @@ class PositionDeleteWriter::Impl {

WriteResult result;
result.data_files.push_back(std::move(data_file));
result.referenced_data_files.assign(referenced_paths_.begin(),
referenced_paths_.end());
return result;
}

Expand All @@ -173,15 +231,29 @@ class PositionDeleteWriter::Impl {
delete_schema_(std::move(delete_schema)),
writer_(std::move(writer)) {}

Status FlushBuffer() {
ArrowSchema arrow_schema;
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*delete_schema_, &arrow_schema));
internal::ArrowSchemaGuard schema_guard(&arrow_schema);
Status InitSchema() {
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*delete_schema_, &arrow_schema_));
ArrowError error;
// The delete schema never changes, so the view is initialized once here and
// merely rebound to each incoming batch in Write.
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
ArrowArrayViewInitFromSchema(&array_view_, &arrow_schema_, &error), error);
return {};
}

/// \brief Undo the referenced paths added by the batch that failed to write.
void RollbackPendingReferences() {
for (auto it : pending_references_) {
referenced_paths_.erase(it);
}
pending_references_.clear();
}

Status FlushBuffer() {
ArrowArray array;
ArrowError error;
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
ArrowArrayInitFromSchema(&array, &arrow_schema, &error), error);
ArrowArrayInitFromSchema(&array, &arrow_schema_, &error), error);
internal::ArrowArrayGuard array_guard(&array);
ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayStartAppending(&array));

Expand All @@ -205,13 +277,24 @@ class PositionDeleteWriter::Impl {
return {};
}

// Transparent comparator so that paths arriving as string_view can be looked up
// without first materializing a std::string.
using ReferencedPaths = std::set<std::string, std::less<>>;

PositionDeleteWriterOptions options_;
std::shared_ptr<Schema> delete_schema_;
std::unique_ptr<Writer> writer_;
// The immutable delete schema in Arrow form, paired with the view bound to it.
// Declared before the view and released after it in the destructor.
ArrowSchema arrow_schema_{};
ArrowArrayView array_view_{};
bool closed_ = false;
std::vector<std::string> buffered_paths_;
std::vector<int64_t> buffered_positions_;
std::set<std::string> referenced_paths_;
ReferencedPaths referenced_paths_;
// Iterators of the entries added to referenced_paths_ by the batch in flight, so
// that they can be removed again if that batch is rejected.
std::vector<ReferencedPaths::iterator> pending_references_;
};

PositionDeleteWriter::PositionDeleteWriter(std::unique_ptr<Impl> impl)
Expand Down
211 changes: 201 additions & 10 deletions src/iceberg/test/data_writer_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
#include <gtest/gtest.h>

#include "iceberg/arrow/arrow_io_internal.h"
#include "iceberg/arrow_c_data_guard_internal.h"
#include "iceberg/avro/avro_register.h"
#include "iceberg/data/equality_delete_writer.h"
#include "iceberg/data/position_delete_writer.h"
Expand All @@ -45,7 +46,9 @@

namespace iceberg {

using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::UnorderedElementsAre;

class DataWriterTest : public ::testing::Test {
protected:
Expand Down Expand Up @@ -344,18 +347,12 @@ class PositionDeleteWriterTest : public DataWriterTest {
};
}

std::shared_ptr<::arrow::Array> CreatePositionDeleteData() {
std::shared_ptr<::arrow::Array> CreatePositionDeleteData(
std::string_view json =
R"([["data_file_1.parquet", 0], ["data_file_1.parquet", 5], ["data_file_1.parquet", 10]])") {
auto delete_schema = std::make_shared<Schema>(std::vector<SchemaField>{
MetadataColumns::kDeleteFilePath, MetadataColumns::kDeleteFilePos});

ArrowSchema arrow_c_schema;
ICEBERG_THROW_NOT_OK(ToArrowSchema(*delete_schema, &arrow_c_schema));
auto arrow_type = ::arrow::ImportType(&arrow_c_schema).ValueOrDie();

return ::arrow::json::ArrayFromJSONString(
::arrow::struct_(arrow_type->fields()),
R"([["data_file_1.parquet", 0], ["data_file_1.parquet", 5], ["data_file_1.parquet", 10]])")
.ValueOrDie();
return CreateArray(*delete_schema, json);
}
};

Expand Down Expand Up @@ -458,6 +455,200 @@ TEST_F(PositionDeleteWriterTest, WriteBatchData) {
const auto& data_file = metadata_result.value().data_files[0];
EXPECT_EQ(data_file->content, DataFile::Content::kPositionDeletes);
EXPECT_GT(data_file->file_size_in_bytes, 0);
ASSERT_TRUE(data_file->referenced_data_file.has_value());
EXPECT_EQ(data_file->referenced_data_file.value(), "data_file_1.parquet");
// Bounds for delete metadata columns are kept when referencing a single file.
EXPECT_TRUE(data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
EXPECT_TRUE(data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
EXPECT_TRUE(data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
EXPECT_TRUE(data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
}

TEST_F(PositionDeleteWriterTest, WriteBatchRejectsSlicedData) {
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

auto test_data = CreatePositionDeleteData(
R"([["data_file_1.parquet", 0], ["data_file_1.parquet", 5]])");
auto sliced = test_data->Slice(1, 1);
ArrowArray arrow_array;
ASSERT_TRUE(::arrow::ExportArray(*sliced, &arrow_array).ok());

auto result = writer->Write(&arrow_array);
EXPECT_EQ(arrow_array.release, nullptr);
internal::ArrowArrayGuard array_guard(&arrow_array);
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
EXPECT_THAT(
result,
HasErrorMessage("Position delete data with a non-zero offset is not supported"));
}

TEST_F(PositionDeleteWriterTest, FailedBatchWriteDoesNotTrackReferencedFiles) {
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

auto good_data = CreatePositionDeleteData(R"([["data_file_1.parquet", 0]])");
ArrowArray good_array;
ASSERT_TRUE(::arrow::ExportArray(*good_data, &good_array).ok());
ASSERT_THAT(writer->Write(&good_array), IsOk());

// The batch references a valid path before the null path rejects it, and none of
// its paths may end up in the metadata.
auto bad_data =
CreatePositionDeleteData(R"([["data_file_bad.parquet", 1], [null, 2]])");
ArrowArray bad_array;
ASSERT_TRUE(::arrow::ExportArray(*bad_data, &bad_array).ok());
internal::ArrowArrayGuard bad_array_guard(&bad_array);
ASSERT_THAT(writer->Write(&bad_array), IsError(ErrorKind::kInvalidArrowData));

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks odd to me because we continue to use a failed writer which should not happen in production. And this does actually verify the case name FailedBatchWriteDoesNotTrackReferencedFiles.


ASSERT_THAT(writer->Close(), IsOk());

auto metadata_result = writer->Metadata();
ASSERT_THAT(metadata_result, IsOk());

const auto& write_result = metadata_result.value();
const auto& data_file = write_result.data_files[0];
ASSERT_TRUE(data_file->referenced_data_file.has_value());
EXPECT_EQ(data_file->referenced_data_file.value(), "data_file_1.parquet");
EXPECT_THAT(write_result.referenced_data_files, ElementsAre("data_file_1.parquet"));
}

TEST_F(PositionDeleteWriterTest, WriteBatchDataForMultipleFiles) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is one batch containing two paths, not multiple successful Write calls. A bug that replaces referenced_paths_ instead of unioning across batches would still pass. Add two successful batches with disjoint paths.

auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

// Disjoint paths across two successful batches must be unioned, not replaced.
auto first_data = CreatePositionDeleteData(R"([["data_file_1.parquet", 0]])");
ArrowArray first_array;
ASSERT_TRUE(::arrow::ExportArray(*first_data, &first_array).ok());
ASSERT_THAT(writer->Write(&first_array), IsOk());

auto second_data = CreatePositionDeleteData(R"([["data_file_2.parquet", 5]])");
ArrowArray second_array;
ASSERT_TRUE(::arrow::ExportArray(*second_data, &second_array).ok());
ASSERT_THAT(writer->Write(&second_array), IsOk());

ASSERT_THAT(writer->Close(), IsOk());

auto metadata_result = writer->Metadata();
ASSERT_THAT(metadata_result, IsOk());

const auto& write_result = metadata_result.value();
const auto& data_file = write_result.data_files[0];
EXPECT_FALSE(data_file->referenced_data_file.has_value());
EXPECT_THAT(write_result.referenced_data_files,
UnorderedElementsAre("data_file_1.parquet", "data_file_2.parquet"));
EXPECT_FALSE(
data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
EXPECT_FALSE(data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
EXPECT_FALSE(
data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
EXPECT_FALSE(data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
}

TEST_F(PositionDeleteWriterTest, WriteBatchThenDeleteTracksAllReferencedFiles) {
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

auto test_data = CreatePositionDeleteData(R"([["data_file_1.parquet", 0]])");
ArrowArray arrow_array;
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
ASSERT_THAT(writer->Write(&arrow_array), IsOk());
ASSERT_THAT(writer->WriteDelete("data_file_2.parquet", 5), IsOk());
ASSERT_THAT(writer->Close(), IsOk());

auto metadata_result = writer->Metadata();
ASSERT_THAT(metadata_result, IsOk());
const auto& write_result = metadata_result.value();
EXPECT_FALSE(write_result.data_files[0]->referenced_data_file.has_value());
EXPECT_THAT(write_result.referenced_data_files,
UnorderedElementsAre("data_file_1.parquet", "data_file_2.parquet"));
}

TEST_F(PositionDeleteWriterTest, WriteBatchRejectsInvalidInput) {
// A null array.
{
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

auto result = writer->Write(nullptr);
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
EXPECT_THAT(result, HasErrorMessage("Position delete data must not be null"));
}

// A null file path.
{
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

auto test_data = CreatePositionDeleteData(R"([[null, 0]])");
ArrowArray arrow_array;
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());

auto result = writer->Write(&arrow_array);
EXPECT_EQ(arrow_array.release, nullptr);
internal::ArrowArrayGuard array_guard(&arrow_array);
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
EXPECT_THAT(result, HasErrorMessage(
"Position delete file paths must not contain null values"));
}

// A null position.
{
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

auto test_data = CreatePositionDeleteData(R"([["data_file_1.parquet", null]])");
ArrowArray arrow_array;
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());

auto result = writer->Write(&arrow_array);
EXPECT_EQ(arrow_array.release, nullptr);
internal::ArrowArrayGuard array_guard(&arrow_array);
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
EXPECT_THAT(result, HasErrorMessage(
"Position delete positions must not contain null values"));
}

// An empty file path.
{
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

auto test_data = CreatePositionDeleteData(R"([["", 0]])");
ArrowArray arrow_array;
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());

auto result = writer->Write(&arrow_array);
EXPECT_EQ(arrow_array.release, nullptr);
internal::ArrowArrayGuard array_guard(&arrow_array);
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
EXPECT_THAT(result, HasErrorMessage("Position delete file paths must not be empty"));
}
}

TEST_F(PositionDeleteWriterTest, WriteEmptyBatchDoesNotAddReferencedFiles) {
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
ASSERT_THAT(writer_result, IsOk());
auto writer = std::move(writer_result.value());

auto test_data = CreatePositionDeleteData("[]");
ArrowArray arrow_array;
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
ASSERT_THAT(writer->Write(&arrow_array), IsOk());
ASSERT_THAT(writer->Close(), IsOk());

auto metadata_result = writer->Metadata();
ASSERT_THAT(metadata_result, IsOk());
EXPECT_FALSE(metadata_result.value().data_files[0]->referenced_data_file.has_value());
}

TEST_F(PositionDeleteWriterTest, AutoFlushOnThreshold) {
Expand Down
Loading