Skip to content

Commit a1777a5

Browse files
u70b3u70b3
committed
fix: track batch position delete references
1 parent f9e83ff commit a1777a5

2 files changed

Lines changed: 123 additions & 12 deletions

File tree

‎src/iceberg/data/position_delete_writer.cc‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,10 +62,40 @@ class PositionDeleteWriter::Impl {
6262
}
6363

6464
Status Write(ArrowArray* data) {
65+
ICEBERG_PRECHECK(data != nullptr, "Position delete data must not be null");
6566
ICEBERG_PRECHECK(buffered_paths_.empty(),
6667
"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);
68+
69+
ArrowSchema arrow_schema;
70+
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*delete_schema_, &arrow_schema));
71+
internal::ArrowSchemaGuard schema_guard(&arrow_schema);
72+
73+
ArrowArrayView array_view;
74+
internal::ArrowArrayViewGuard view_guard(&array_view);
75+
ArrowError error;
76+
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
77+
ArrowArrayViewInitFromSchema(&array_view, &arrow_schema, &error), error);
78+
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
79+
ArrowArrayViewSetArray(&array_view, data, &error), error);
80+
81+
const auto* path_view = array_view.children[0];
82+
if (ArrowArrayViewComputeNullCount(path_view) != 0) {
83+
return InvalidArrowData("Position delete file paths must not contain null values");
84+
}
85+
86+
std::set<std::string> pending_paths;
87+
for (int64_t i = 0; i < data->length; ++i) {
88+
auto path = ArrowArrayViewGetStringUnsafe(path_view, i);
89+
if (path.size_bytes == 0) {
90+
pending_paths.emplace();
91+
} else {
92+
pending_paths.emplace(path.data, static_cast<size_t>(path.size_bytes));
93+
}
94+
}
95+
96+
ICEBERG_RETURN_UNEXPECTED(writer_->Write(data));
97+
referenced_paths_.merge(pending_paths);
98+
return {};
6999
}
70100

71101
Status WriteDelete(std::string_view file_path, int64_t pos) {

‎src/iceberg/test/data_writer_test.cc‎

Lines changed: 91 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"
@@ -344,18 +345,12 @@ class PositionDeleteWriterTest : public DataWriterTest {
344345
};
345346
}
346347

347-
std::shared_ptr<::arrow::Array> CreatePositionDeleteData() {
348+
std::shared_ptr<::arrow::Array> CreatePositionDeleteData(
349+
std::string_view json =
350+
R"([["data_file_1.parquet", 0], ["data_file_1.parquet", 5], ["data_file_1.parquet", 10]])") {
348351
auto delete_schema = std::make_shared<Schema>(std::vector<SchemaField>{
349352
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();
353+
return CreateArray(*delete_schema, json);
359354
}
360355
};
361356

@@ -458,6 +453,92 @@ TEST_F(PositionDeleteWriterTest, WriteBatchData) {
458453
const auto& data_file = metadata_result.value().data_files[0];
459454
EXPECT_EQ(data_file->content, DataFile::Content::kPositionDeletes);
460455
EXPECT_GT(data_file->file_size_in_bytes, 0);
456+
ASSERT_TRUE(data_file->referenced_data_file.has_value());
457+
EXPECT_EQ(data_file->referenced_data_file.value(), "data_file_1.parquet");
458+
}
459+
460+
TEST_F(PositionDeleteWriterTest, WriteBatchDataForMultipleFiles) {
461+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
462+
ASSERT_THAT(writer_result, IsOk());
463+
auto writer = std::move(writer_result.value());
464+
465+
auto test_data = CreatePositionDeleteData(
466+
R"([["data_file_1.parquet", 0], ["data_file_2.parquet", 5]])");
467+
ArrowArray arrow_array;
468+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
469+
ASSERT_THAT(writer->Write(&arrow_array), IsOk());
470+
ASSERT_THAT(writer->Close(), IsOk());
471+
472+
auto metadata_result = writer->Metadata();
473+
ASSERT_THAT(metadata_result, IsOk());
474+
475+
const auto& data_file = metadata_result.value().data_files[0];
476+
EXPECT_FALSE(data_file->referenced_data_file.has_value());
477+
EXPECT_FALSE(
478+
data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
479+
EXPECT_FALSE(data_file->lower_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
480+
EXPECT_FALSE(
481+
data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePathColumnId));
482+
EXPECT_FALSE(data_file->upper_bounds.contains(MetadataColumns::kDeleteFilePosColumnId));
483+
}
484+
485+
TEST_F(PositionDeleteWriterTest, WriteBatchThenDeleteTracksAllReferencedFiles) {
486+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
487+
ASSERT_THAT(writer_result, IsOk());
488+
auto writer = std::move(writer_result.value());
489+
490+
auto test_data = CreatePositionDeleteData(R"([["data_file_1.parquet", 0]])");
491+
ArrowArray arrow_array;
492+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
493+
ASSERT_THAT(writer->Write(&arrow_array), IsOk());
494+
ASSERT_THAT(writer->WriteDelete("data_file_2.parquet", 5), IsOk());
495+
ASSERT_THAT(writer->Close(), IsOk());
496+
497+
auto metadata_result = writer->Metadata();
498+
ASSERT_THAT(metadata_result, IsOk());
499+
EXPECT_FALSE(metadata_result.value().data_files[0]->referenced_data_file.has_value());
500+
}
501+
502+
TEST_F(PositionDeleteWriterTest, WriteBatchRejectsNullFilePath) {
503+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
504+
ASSERT_THAT(writer_result, IsOk());
505+
auto writer = std::move(writer_result.value());
506+
507+
auto test_data = CreatePositionDeleteData(R"([[null, 0]])");
508+
ArrowArray arrow_array;
509+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
510+
internal::ArrowArrayGuard array_guard(&arrow_array);
511+
512+
auto result = writer->Write(&arrow_array);
513+
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
514+
EXPECT_THAT(result,
515+
HasErrorMessage("Position delete file paths must not contain null values"));
516+
}
517+
518+
TEST_F(PositionDeleteWriterTest, WriteBatchRejectsNullData) {
519+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
520+
ASSERT_THAT(writer_result, IsOk());
521+
auto writer = std::move(writer_result.value());
522+
523+
auto result = writer->Write(nullptr);
524+
ASSERT_THAT(result, IsError(ErrorKind::kInvalidArgument));
525+
EXPECT_THAT(result, HasErrorMessage("Position delete data must not be null"));
526+
}
527+
528+
TEST_F(PositionDeleteWriterTest, WriteEmptyBatchDoesNotAddReferencedFiles) {
529+
auto writer_result = PositionDeleteWriter::Make(MakeDeleteOptions());
530+
ASSERT_THAT(writer_result, IsOk());
531+
auto writer = std::move(writer_result.value());
532+
533+
auto test_data = CreatePositionDeleteData("[]");
534+
ArrowArray arrow_array;
535+
ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
536+
ASSERT_THAT(writer->Write(&arrow_array), IsOk());
537+
ASSERT_THAT(writer->Close(), IsOk());
538+
539+
auto metadata_result = writer->Metadata();
540+
ASSERT_THAT(metadata_result, IsOk());
541+
EXPECT_FALSE(metadata_result.value().data_files[0]->referenced_data_file.has_value());
461542
}
462543

463544
TEST_F(PositionDeleteWriterTest, AutoFlushOnThreshold) {

0 commit comments

Comments
 (0)