Skip to content

Commit 307adb1

Browse files
u70b3u70b3
authored
fix: track batch position delete references (#842)
1 parent 4e4e10a commit 307adb1

2 files changed

Lines changed: 293 additions & 19 deletions

File tree

‎src/iceberg/data/position_delete_writer.cc‎

Lines changed: 92 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,11 @@
1919

2020
#include "iceberg/data/position_delete_writer.h"
2121

22+
#include <functional>
2223
#include <map>
2324
#include <set>
25+
#include <string>
26+
#include <string_view>
2427
#include <utility>
2528
#include <vector>
2629

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

60-
return std::unique_ptr<Impl>(
63+
auto impl = std::unique_ptr<Impl>(
6164
new Impl(std::move(options), std::move(delete_schema), std::move(writer)));
65+
ICEBERG_RETURN_UNEXPECTED(impl->InitSchema());
66+
return impl;
67+
}
68+
69+
~Impl() {
70+
ArrowArrayViewReset(&array_view_);
71+
if (arrow_schema_.release != nullptr) {
72+
ArrowSchemaRelease(&arrow_schema_);
73+
}
6274
}
6375

6476
Status Write(ArrowArray* data) {
77+
ICEBERG_PRECHECK(data != nullptr, "Position delete data must not be null");
78+
internal::ArrowArrayGuard data_guard(data);
79+
ICEBERG_PRECHECK(data->offset == 0,
80+
"Position delete data with a non-zero offset is not supported");
6581
ICEBERG_PRECHECK(buffered_paths_.empty(),
6682
"Cannot write batch data when there are buffered deletes.");
67-
// TODO(anyone): Extract file paths from ArrowArray to update referenced_paths_.
68-
return writer_->Write(data);
83+
84+
ArrowError error;
85+
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
86+
ArrowArrayViewSetArray(&array_view_, data, &error), error);
87+
88+
const auto* path_view = array_view_.children[0];
89+
const auto* pos_view = array_view_.children[1];
90+
91+
// A batch usually references files that earlier batches already referenced, so
92+
// record paths optimistically: a known path costs a lookup instead of a scratch
93+
// allocation. Entries added by this batch are rolled back unless the write
94+
// succeeds, so a rejected batch still leaves no trace in the metadata.
95+
pending_references_.clear();
96+
for (int64_t i = 0; i < data->length; ++i) {
97+
if (ArrowArrayViewIsNull(path_view, i)) {
98+
RollbackPendingReferences();
99+
return InvalidArrowData(
100+
"Position delete file paths must not contain null values");
101+
}
102+
if (ArrowArrayViewIsNull(pos_view, i)) {
103+
RollbackPendingReferences();
104+
return InvalidArrowData("Position delete positions must not contain null values");
105+
}
106+
auto path = ArrowArrayViewGetStringUnsafe(path_view, i);
107+
if (path.size_bytes == 0) {
108+
RollbackPendingReferences();
109+
return InvalidArrowData("Position delete file paths must not be empty");
110+
}
111+
std::string_view file_path(path.data, static_cast<size_t>(path.size_bytes));
112+
if (!referenced_paths_.contains(file_path)) {
113+
pending_references_.push_back(
114+
referenced_paths_.insert(std::string(file_path)).first);
115+
}
116+
}
117+
118+
Status status = writer_->Write(data);
119+
if (!status) {
120+
RollbackPendingReferences();
121+
return status;
122+
}
123+
pending_references_.clear();
124+
return {};
69125
}
70126

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

164220
WriteResult result;
165221
result.data_files.push_back(std::move(data_file));
222+
result.referenced_data_files.assign(referenced_paths_.begin(),
223+
referenced_paths_.end());
166224
return result;
167225
}
168226

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

176-
Status FlushBuffer() {
177-
ArrowSchema arrow_schema;
178-
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*delete_schema_, &arrow_schema));
179-
internal::ArrowSchemaGuard schema_guard(&arrow_schema);
234+
Status InitSchema() {
235+
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*delete_schema_, &arrow_schema_));
236+
ArrowError error;
237+
// The delete schema never changes, so the view is initialized once here and
238+
// merely rebound to each incoming batch in Write.
239+
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
240+
ArrowArrayViewInitFromSchema(&array_view_, &arrow_schema_, &error), error);
241+
return {};
242+
}
180243

244+
/// \brief Undo the referenced paths added by the batch that failed to write.
245+
void RollbackPendingReferences() {
246+
for (auto it : pending_references_) {
247+
referenced_paths_.erase(it);
248+
}
249+
pending_references_.clear();
250+
}
251+
252+
Status FlushBuffer() {
181253
ArrowArray array;
182254
ArrowError error;
183255
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
184-
ArrowArrayInitFromSchema(&array, &arrow_schema, &error), error);
256+
ArrowArrayInitFromSchema(&array, &arrow_schema_, &error), error);
185257
internal::ArrowArrayGuard array_guard(&array);
186258
ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayStartAppending(&array));
187259

@@ -205,13 +277,24 @@ class PositionDeleteWriter::Impl {
205277
return {};
206278
}
207279

280+
// Transparent comparator so that paths arriving as string_view can be looked up
281+
// without first materializing a std::string.
282+
using ReferencedPaths = std::set<std::string, std::less<>>;
283+
208284
PositionDeleteWriterOptions options_;
209285
std::shared_ptr<Schema> delete_schema_;
210286
std::unique_ptr<Writer> writer_;
287+
// The immutable delete schema in Arrow form, paired with the view bound to it.
288+
// Declared before the view and released after it in the destructor.
289+
ArrowSchema arrow_schema_{};
290+
ArrowArrayView array_view_{};
211291
bool closed_ = false;
212292
std::vector<std::string> buffered_paths_;
213293
std::vector<int64_t> buffered_positions_;
214-
std::set<std::string> referenced_paths_;
294+
ReferencedPaths referenced_paths_;
295+
// Iterators of the entries added to referenced_paths_ by the batch in flight, so
296+
// that they can be removed again if that batch is rejected.
297+
std::vector<ReferencedPaths::iterator> pending_references_;
215298
};
216299

217300
PositionDeleteWriter::PositionDeleteWriter(std::unique_ptr<Impl> impl)

‎src/iceberg/test/data_writer_test.cc‎

Lines changed: 201 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
#include <gtest/gtest.h>
2727

2828
#include "iceberg/arrow/arrow_io_internal.h"
29+
#include "iceberg/arrow_c_data_guard_internal.h"
2930
#include "iceberg/avro/avro_register.h"
3031
#include "iceberg/data/equality_delete_writer.h"
3132
#include "iceberg/data/position_delete_writer.h"
@@ -45,7 +46,9 @@
4546

4647
namespace iceberg {
4748

49+
using ::testing::ElementsAre;
4850
using ::testing::HasSubstr;
51+
using ::testing::UnorderedElementsAre;
4952

5053
class DataWriterTest : public ::testing::Test {
5154
protected:
@@ -344,18 +347,12 @@ class PositionDeleteWriterTest : public DataWriterTest {
344347
};
345348
}
346349

347-
std::shared_ptr<::arrow::Array> CreatePositionDeleteData() {
350+
std::shared_ptr<::arrow::Array> CreatePositionDeleteData(
351+
std::string_view json =
352+
R"([["data_file_1.parquet", 0], ["data_file_1.parquet", 5], ["data_file_1.parquet", 10]])") {
348353
auto delete_schema = std::make_shared<Schema>(std::vector<SchemaField>{
349354
MetadataColumns::kDeleteFilePath, MetadataColumns::kDeleteFilePos});
350-
351-
ArrowSchema arrow_c_schema;
352-
ICEBERG_THROW_NOT_OK(ToArrowSchema(*delete_schema, &arrow_c_schema));
353-
auto arrow_type = ::arrow::ImportType(&arrow_c_schema).ValueOrDie();
354-
355-
return ::arrow::json::ArrayFromJSONString(
356-
::arrow::struct_(arrow_type->fields()),
357-
R"([["data_file_1.parquet", 0], ["data_file_1.parquet", 5], ["data_file_1.parquet", 10]])")
358-
.ValueOrDie();
355+
return CreateArray(*delete_schema, json);
359356
}
360357
};
361358

@@ -458,6 +455,200 @@ TEST_F(PositionDeleteWriterTest, WriteBatchData) {
458455
const auto& data_file = metadata_result.value().data_files[0];
459456
EXPECT_EQ(data_file->content, DataFile::Content::kPositionDeletes);
460457
EXPECT_GT(data_file->file_size_in_bytes, 0);
458+
ASSERT_TRUE(data_file->referenced_data_file.has_value());
459+
EXPECT_EQ(data_file->referenced_data_file.value(), "data_file_1.parquet");
460+
// Bounds for delete metadata columns are kept when referencing a single file.
461+
EXPECT_TRUE(data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
462+
EXPECT_TRUE(data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
463+
EXPECT_TRUE(data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
464+
EXPECT_TRUE(data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
465+
}
466+
467+
TEST_F(PositionDeleteWriterTest, WriteBatchRejectsSlicedData) {
468+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
469+
ASSERT_THAT(writer_result, IsOk());
470+
auto writer = std::move(writer_result.value());
471+
472+
auto test_data = CreatePositionDeleteData(
473+
R"([["data_file_1.parquet", 0], ["data_file_1.parquet", 5]])");
474+
auto sliced = test_data->Slice(1, 1);
475+
ArrowArray arrow_array;
476+
ASSERT_TRUE(::arrow::ExportArray(*sliced, &arrow_array).ok());
477+
478+
auto result = writer->Write(&arrow_array);
479+
EXPECT_EQ(arrow_array.release, nullptr);
480+
internal::ArrowArrayGuard array_guard(&arrow_array);
481+
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
482+
EXPECT_THAT(
483+
result,
484+
HasErrorMessage("Position delete data with a non-zero offset is not supported"));
485+
}
486+
487+
TEST_F(PositionDeleteWriterTest, FailedBatchWriteDoesNotTrackReferencedFiles) {
488+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
489+
ASSERT_THAT(writer_result, IsOk());
490+
auto writer = std::move(writer_result.value());
491+
492+
auto good_data = CreatePositionDeleteData(R"([["data_file_1.parquet", 0]])");
493+
ArrowArray good_array;
494+
ASSERT_TRUE(::arrow::ExportArray(*good_data, &good_array).ok());
495+
ASSERT_THAT(writer->Write(&good_array), IsOk());
496+
497+
// The batch references a valid path before the null path rejects it, and none of
498+
// its paths may end up in the metadata.
499+
auto bad_data =
500+
CreatePositionDeleteData(R"([["data_file_bad.parquet", 1], [null, 2]])");
501+
ArrowArray bad_array;
502+
ASSERT_TRUE(::arrow::ExportArray(*bad_data, &bad_array).ok());
503+
internal::ArrowArrayGuard bad_array_guard(&bad_array);
504+
ASSERT_THAT(writer->Write(&bad_array), IsError(ErrorKind::kInvalidArrowData));
505+
506+
ASSERT_THAT(writer->Close(), IsOk());
507+
508+
auto metadata_result = writer->Metadata();
509+
ASSERT_THAT(metadata_result, IsOk());
510+
511+
const auto& write_result = metadata_result.value();
512+
const auto& data_file = write_result.data_files[0];
513+
ASSERT_TRUE(data_file->referenced_data_file.has_value());
514+
EXPECT_EQ(data_file->referenced_data_file.value(), "data_file_1.parquet");
515+
EXPECT_THAT(write_result.referenced_data_files, ElementsAre("data_file_1.parquet"));
516+
}
517+
518+
TEST_F(PositionDeleteWriterTest, WriteBatchDataForMultipleFiles) {
519+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
520+
ASSERT_THAT(writer_result, IsOk());
521+
auto writer = std::move(writer_result.value());
522+
523+
// Disjoint paths across two successful batches must be unioned, not replaced.
524+
auto first_data = CreatePositionDeleteData(R"([["data_file_1.parquet", 0]])");
525+
ArrowArray first_array;
526+
ASSERT_TRUE(::arrow::ExportArray(*first_data, &first_array).ok());
527+
ASSERT_THAT(writer->Write(&first_array), IsOk());
528+
529+
auto second_data = CreatePositionDeleteData(R"([["data_file_2.parquet", 5]])");
530+
ArrowArray second_array;
531+
ASSERT_TRUE(::arrow::ExportArray(*second_data, &second_array).ok());
532+
ASSERT_THAT(writer->Write(&second_array), IsOk());
533+
534+
ASSERT_THAT(writer->Close(), IsOk());
535+
536+
auto metadata_result = writer->Metadata();
537+
ASSERT_THAT(metadata_result, IsOk());
538+
539+
const auto& write_result = metadata_result.value();
540+
const auto& data_file = write_result.data_files[0];
541+
EXPECT_FALSE(data_file->referenced_data_file.has_value());
542+
EXPECT_THAT(write_result.referenced_data_files,
543+
UnorderedElementsAre("data_file_1.parquet", "data_file_2.parquet"));
544+
EXPECT_FALSE(
545+
data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
546+
EXPECT_FALSE(data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
547+
EXPECT_FALSE(
548+
data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
549+
EXPECT_FALSE(data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
550+
}
551+
552+
TEST_F(PositionDeleteWriterTest, WriteBatchThenDeleteTracksAllReferencedFiles) {
553+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
554+
ASSERT_THAT(writer_result, IsOk());
555+
auto writer = std::move(writer_result.value());
556+
557+
auto test_data = CreatePositionDeleteData(R"([["data_file_1.parquet", 0]])");
558+
ArrowArray arrow_array;
559+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
560+
ASSERT_THAT(writer->Write(&arrow_array), IsOk());
561+
ASSERT_THAT(writer->WriteDelete("data_file_2.parquet", 5), IsOk());
562+
ASSERT_THAT(writer->Close(), IsOk());
563+
564+
auto metadata_result = writer->Metadata();
565+
ASSERT_THAT(metadata_result, IsOk());
566+
const auto& write_result = metadata_result.value();
567+
EXPECT_FALSE(write_result.data_files[0]->referenced_data_file.has_value());
568+
EXPECT_THAT(write_result.referenced_data_files,
569+
UnorderedElementsAre("data_file_1.parquet", "data_file_2.parquet"));
570+
}
571+
572+
TEST_F(PositionDeleteWriterTest, WriteBatchRejectsInvalidInput) {
573+
// A null array.
574+
{
575+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
576+
ASSERT_THAT(writer_result, IsOk());
577+
auto writer = std::move(writer_result.value());
578+
579+
auto result = writer->Write(nullptr);
580+
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
581+
EXPECT_THAT(result, HasErrorMessage("Position delete data must not be null"));
582+
}
583+
584+
// A null file path.
585+
{
586+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
587+
ASSERT_THAT(writer_result, IsOk());
588+
auto writer = std::move(writer_result.value());
589+
590+
auto test_data = CreatePositionDeleteData(R"([[null, 0]])");
591+
ArrowArray arrow_array;
592+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
593+
594+
auto result = writer->Write(&arrow_array);
595+
EXPECT_EQ(arrow_array.release, nullptr);
596+
internal::ArrowArrayGuard array_guard(&arrow_array);
597+
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
598+
EXPECT_THAT(result, HasErrorMessage(
599+
"Position delete file paths must not contain null values"));
600+
}
601+
602+
// A null position.
603+
{
604+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
605+
ASSERT_THAT(writer_result, IsOk());
606+
auto writer = std::move(writer_result.value());
607+
608+
auto test_data = CreatePositionDeleteData(R"([["data_file_1.parquet", null]])");
609+
ArrowArray arrow_array;
610+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
611+
612+
auto result = writer->Write(&arrow_array);
613+
EXPECT_EQ(arrow_array.release, nullptr);
614+
internal::ArrowArrayGuard array_guard(&arrow_array);
615+
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
616+
EXPECT_THAT(result, HasErrorMessage(
617+
"Position delete positions must not contain null values"));
618+
}
619+
620+
// An empty file path.
621+
{
622+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
623+
ASSERT_THAT(writer_result, IsOk());
624+
auto writer = std::move(writer_result.value());
625+
626+
auto test_data = CreatePositionDeleteData(R"([["", 0]])");
627+
ArrowArray arrow_array;
628+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
629+
630+
auto result = writer->Write(&arrow_array);
631+
EXPECT_EQ(arrow_array.release, nullptr);
632+
internal::ArrowArrayGuard array_guard(&arrow_array);
633+
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
634+
EXPECT_THAT(result, HasErrorMessage("Position delete file paths must not be empty"));
635+
}
636+
}
637+
638+
TEST_F(PositionDeleteWriterTest, WriteEmptyBatchDoesNotAddReferencedFiles) {
639+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
640+
ASSERT_THAT(writer_result, IsOk());
641+
auto writer = std::move(writer_result.value());
642+
643+
auto test_data = CreatePositionDeleteData("[]");
644+
ArrowArray arrow_array;
645+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
646+
ASSERT_THAT(writer->Write(&arrow_array), IsOk());
647+
ASSERT_THAT(writer->Close(), IsOk());
648+
649+
auto metadata_result = writer->Metadata();
650+
ASSERT_THAT(metadata_result, IsOk());
651+
EXPECT_FALSE(metadata_result.value().data_files[0]->referenced_data_file.has_value());
461652
}
462653

463654
TEST_F(PositionDeleteWriterTest, AutoFlushOnThreshold) {

0 commit comments

Comments
 (0)