Skip to content

Commit a899511

Browse files
authored
fix(arrow): make full batch validation configurable (#1043)
1 parent 942e9af commit a899511

6 files changed

Lines changed: 185 additions & 44 deletions

File tree

‎src/iceberg/arrow_c_data_util.cc‎

Lines changed: 28 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,8 @@ ProjectionContext::ProjectionContext(ProjectionContext&& other) noexcept
238238
input_arrow_schema_(other.input_arrow_schema_),
239239
output_arrow_schema_(other.output_arrow_schema_),
240240
project_batch_function_(std::exchange(other.project_batch_function_, nullptr)),
241-
project_batch_state_(std::move(other.project_batch_state_)) {
241+
project_batch_state_(std::move(other.project_batch_state_)),
242+
validate_batch_(other.validate_batch_) {
242243
other.input_arrow_schema_.release = nullptr;
243244
other.output_arrow_schema_.release = nullptr;
244245
}
@@ -264,6 +265,7 @@ ProjectionContext& ProjectionContext::operator=(ProjectionContext&& other) noexc
264265
other.output_arrow_schema_.release = nullptr;
265266
project_batch_function_ = std::exchange(other.project_batch_function_, nullptr);
266267
project_batch_state_ = std::move(other.project_batch_state_);
268+
validate_batch_ = other.validate_batch_;
267269
return *this;
268270
}
269271

@@ -278,7 +280,7 @@ ProjectionContext::~ProjectionContext() {
278280

279281
Result<ProjectionContext> ProjectionContext::Make(
280282
const Schema& input_schema, const Schema& output_schema,
281-
ProjectionContext::ProjectBatchFunction project_batch_function) {
283+
ProjectionContext::ProjectBatchFunction project_batch_function, bool validate_batch) {
282284
ICEBERG_ASSIGN_OR_RAISE(
283285
auto selected_field_indices,
284286
BuildSelectedFieldIndices(input_schema.fields(), output_schema.fields()));
@@ -290,6 +292,7 @@ Result<ProjectionContext> ProjectionContext::Make(
290292
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(input_schema, &context.input_arrow_schema_));
291293
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(output_schema, &context.output_arrow_schema_));
292294
context.project_batch_function_ = project_batch_function;
295+
context.validate_batch_ = validate_batch;
293296

294297
return context;
295298
}
@@ -319,6 +322,8 @@ ProjectionContext::ProjectBatchState& ProjectionContext::project_batch_state() {
319322
return project_batch_state_;
320323
}
321324

325+
bool ProjectionContext::validate_batch() const { return validate_batch_; }
326+
322327
void ProjectionContext::RegisterProjectBatchFunction(
323328
ProjectionContext::ProjectBatchFunction project_batch_function) {
324329
ICEBERG_DCHECK(project_batch_function != nullptr,
@@ -341,21 +346,13 @@ auto ProjectionContext::ResolveProjectBatchFunction()
341346

342347
namespace {
343348

344-
Result<ArrowArray> ProjectBatchNanoarrow(
345-
const ArrowSchema& input_arrow_schema, const ArrowArray& input_batch,
346-
std::span<const int32_t> row_indices, const ArrowSchema& output_arrow_schema,
347-
std::span<const int32_t> selected_field_indices) {
348-
ArrowArrayView input_view;
349+
Result<ArrowArray> ProjectBatchNanoarrow(const ArrowSchema& input_arrow_schema,
350+
const ArrowArray& input_batch,
351+
std::span<const int32_t> row_indices,
352+
const ArrowSchema& output_arrow_schema,
353+
std::span<const int32_t> selected_field_indices,
354+
const ArrowArrayView& input_view) {
349355
ArrowError error;
350-
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
351-
ArrowArrayViewInitFromSchema(&input_view, &input_arrow_schema, &error), error);
352-
internal::ArrowArrayViewGuard input_view_guard(&input_view);
353-
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
354-
ArrowArrayViewSetArray(&input_view, &input_batch, &error), error);
355-
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
356-
ArrowArrayViewValidate(&input_view, NANOARROW_VALIDATION_LEVEL_DEFAULT, &error),
357-
error);
358-
359356
ArrowArray output_array;
360357
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
361358
ArrowArrayInitFromSchema(&output_array, &output_arrow_schema, &error), error);
@@ -393,6 +390,20 @@ Result<ArrowArray> ProjectBatch(ArrowArray* input_batch,
393390
ICEBERG_PRECHECK(input_batch != nullptr, "input_batch must not be null");
394391
internal::ArrowArrayGuard input_batch_guard(input_batch);
395392

393+
ArrowArrayView input_view;
394+
ArrowError error;
395+
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
396+
ArrowArrayViewInitFromSchema(&input_view, &projection.input_arrow_schema(), &error),
397+
error);
398+
internal::ArrowArrayViewGuard input_view_guard(&input_view);
399+
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
400+
ArrowArrayViewSetArray(&input_view, input_batch, &error), error);
401+
if (projection.validate_batch()) {
402+
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
403+
ArrowArrayViewValidate(&input_view, NANOARROW_VALIDATION_LEVEL_FULL, &error),
404+
error);
405+
}
406+
396407
auto project_batch_function = projection.project_batch_function();
397408
if (project_batch_function != nullptr) {
398409
// ProjectBatch owns input_batch. Arrow-backed implementations import it and clear
@@ -402,7 +413,7 @@ Result<ArrowArray> ProjectBatch(ArrowArray* input_batch,
402413

403414
return ProjectBatchNanoarrow(projection.input_arrow_schema(), *input_batch, row_indices,
404415
projection.output_arrow_schema(),
405-
projection.selected_field_indices());
416+
projection.selected_field_indices(), input_view);
406417
}
407418

408419
namespace {

‎src/iceberg/arrow_c_data_util_internal.h‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -68,13 +68,16 @@ class ICEBERG_EXPORT ProjectionContext {
6868
/// \param project_batch_function Optional implementation returned by
6969
/// ProjectionContext::ResolveProjectBatchFunction, or nullptr to use the nanoarrow
7070
/// path.
71+
/// \param validate_batch Whether to perform full content validation. Disable only for
72+
/// trusted input.
7173
/// \note It validates that output_schema selects or reorders complete top-level fields
7274
/// by field id. Nested pruning and type changes are rejected. The input_schema and
7375
/// output_schema passed to Make must outlive the context. ProjectBatch may lazily
7476
/// initialize backend cache; do not share one context across concurrent calls.
7577
static Result<ProjectionContext> Make(const Schema& input_schema,
7678
const Schema& output_schema,
77-
ProjectBatchFunction project_batch_function);
79+
ProjectBatchFunction project_batch_function,
80+
bool validate_batch = true);
7881

7982
ProjectionContext(ProjectionContext&&) noexcept;
8083
ProjectionContext& operator=(ProjectionContext&&) noexcept;
@@ -97,6 +100,8 @@ class ICEBERG_EXPORT ProjectionContext {
97100

98101
ProjectBatchState& project_batch_state();
99102

103+
bool validate_batch() const;
104+
100105
private:
101106
ProjectionContext() = default;
102107

@@ -109,6 +114,7 @@ class ICEBERG_EXPORT ProjectionContext {
109114
ArrowSchema output_arrow_schema_{};
110115
ProjectBatchFunction project_batch_function_ = nullptr;
111116
ProjectBatchState project_batch_state_;
117+
bool validate_batch_ = true;
112118
};
113119

114120
/// \brief Concept for sources that can be wrapped as ArrowArrayStreams.
@@ -228,10 +234,11 @@ ICEBERG_EXPORT Result<ArrowArrayStream> MakeArrowArrayStream(
228234
/// \brief Project selected rows from a batch into complete top-level fields.
229235
///
230236
/// `input_batch` is consumed by this function. If the projection carries a registered
231-
/// implementation, the call is delegated to it; otherwise the built-in nanoarrow
232-
/// implementation is used. The projection must have been created for the stable schema
233-
/// pair that describes the input batch and requested output. This function does not
234-
/// revalidate schema compatibility on each batch.
237+
/// implementation, the call is delegated to it after validation; otherwise the built-in
238+
/// nanoarrow implementation is used. The projection must have been created for the stable
239+
/// schema pair that describes the input batch and requested output. This function does
240+
/// not revalidate schema compatibility on each batch. Input arrays are validated
241+
/// according to the context policy before either backend runs.
235242
///
236243
/// \param input_batch Owned Arrow C Data batch to project.
237244
/// \param row_indices Zero-based row positions to copy from `input_batch`.

‎src/iceberg/data/file_scan_task_reader.cc‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -213,9 +213,11 @@ class FileScanTaskReader::Impl {
213213
auto required_schema = delete_filter->RequiredSchema();
214214
auto project_batch_function = ProjectionContext::ResolveProjectBatchFunction();
215215
// ProjectionContext borrows schemas that are kept in MergeOnReadStreamSource.
216-
ICEBERG_ASSIGN_OR_RAISE(auto projection_context,
217-
ProjectionContext::Make(*required_schema, *projected_schema_,
218-
project_batch_function));
216+
ICEBERG_ASSIGN_OR_RAISE(
217+
auto projection_context,
218+
ProjectionContext::Make(*required_schema, *projected_schema_,
219+
project_batch_function,
220+
properties_.Get(ReaderProperties::kArrowValidateBatch)));
219221

220222
auto options = MakeReaderOptions(task, io_, required_schema, filter, case_sensitive_,
221223
name_mapping_, properties_, data_file->first_row_id,

‎src/iceberg/file_reader.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,9 @@ class ICEBERG_EXPORT ReaderProperties : public ConfigBase<ReaderProperties> {
8080
/// Only the Parquet reader honors this option; other readers ignore it.
8181
/// Default: false (use 32-bit offset list).
8282
inline static Entry<bool> kArrowUseLargeList{"read.arrow.use-large-list", false};
83+
/// \brief Validate Arrow batches before merge-on-read projection.
84+
/// Default: true. Disable only when the reader is trusted to produce valid arrays.
85+
inline static Entry<bool> kArrowValidateBatch{"read.arrow.validate-batch", true};
8386
/// \brief Use footer statistics to prune Parquet row groups.
8487
inline static Entry<bool> kParquetRowGroupFilter{
8588
"read.parquet.row-group-filter.enabled", true};

‎src/iceberg/test/arrow_c_data_util_test.cc‎

Lines changed: 94 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
/// \file arrow_c_data_util_test.cc
2121
/// Verifies ProjectBatch behavior across registered implementations.
2222

23+
#include <array>
2324
#include <cstdint>
2425
#include <memory>
2526
#include <string>
@@ -76,15 +77,15 @@ ProjectionContext::ProjectBatchFunction ArrowComputeFunction() {
7677
std::shared_ptr<::arrow::RecordBatch> RunProjectBatch(
7778
const ::arrow::RecordBatch& batch, const std::vector<int32_t>& alive_indices,
7879
const Schema& required_schema, const Schema& projected_schema,
79-
ProjectionContext::ProjectBatchFunction project_batch_function) {
80+
ProjectionContext::ProjectBatchFunction project_batch_function, bool validate_batch) {
8081
ArrowSchema c_schema;
8182
ArrowArray c_array;
8283
EXPECT_TRUE(::arrow::ExportRecordBatch(batch, &c_array, &c_schema).ok());
8384
ArrowSchemaGuard schema_guard(&c_schema);
8485
ArrowArrayGuard array_guard(&c_array);
8586

86-
auto projection =
87-
ProjectionContext::Make(required_schema, projected_schema, project_batch_function);
87+
auto projection = ProjectionContext::Make(required_schema, projected_schema,
88+
project_batch_function, validate_batch);
8889
EXPECT_THAT(projection, IsOk());
8990

9091
auto result = ProjectBatch(&c_array, alive_indices, projection.value());
@@ -103,22 +104,15 @@ void ExpectProjectBatch(const ::arrow::RecordBatch& batch,
103104
const Schema& required_schema, const Schema& projected_schema,
104105
std::string_view expected_json) {
105106
auto expected = MakeBatch(projected_schema, expected_json);
106-
auto nanoarrow =
107-
RunProjectBatch(batch, alive_indices, required_schema, projected_schema, nullptr);
108-
auto arrow_compute = RunProjectBatch(batch, alive_indices, required_schema,
109-
projected_schema, ArrowComputeFunction());
110-
111-
EXPECT_TRUE(nanoarrow->Equals(*expected)) << "nanoarrow:\n"
112-
<< nanoarrow->ToString() << "expected:\n"
113-
<< expected->ToString();
114-
EXPECT_TRUE(arrow_compute->Equals(*expected))
115-
<< "arrow_compute:\n"
116-
<< arrow_compute->ToString() << "expected:\n"
117-
<< expected->ToString();
118-
EXPECT_TRUE(nanoarrow->Equals(*arrow_compute))
119-
<< "nanoarrow:\n"
120-
<< nanoarrow->ToString() << "arrow_compute:\n"
121-
<< arrow_compute->ToString();
107+
for (bool validate_batch : {true, false}) {
108+
auto nanoarrow = RunProjectBatch(batch, alive_indices, required_schema,
109+
projected_schema, nullptr, validate_batch);
110+
auto arrow_compute =
111+
RunProjectBatch(batch, alive_indices, required_schema, projected_schema,
112+
ArrowComputeFunction(), validate_batch);
113+
EXPECT_TRUE(nanoarrow->Equals(*expected)) << nanoarrow->ToString();
114+
EXPECT_TRUE(arrow_compute->Equals(*expected)) << arrow_compute->ToString();
115+
}
122116
}
123117

124118
std::shared_ptr<Schema> MakeFullSchema() {
@@ -231,6 +225,87 @@ TEST(ProjectBatchTest, ProjectionRejectsNestedPruning) {
231225
EXPECT_THAT(projection, IsError(ErrorKind::kInvalidArgument));
232226
}
233227

228+
TEST(ProjectBatchTest, ValidateBatchControlsNonMonotonicStringOffsets) {
229+
for (bool validate_batch : {true, false}) {
230+
for (auto function :
231+
{ProjectionContext::ProjectBatchFunction{nullptr}, ArrowComputeFunction()}) {
232+
auto schema = std::make_shared<Schema>(
233+
std::vector<SchemaField>{SchemaField::MakeRequired(1, "value", string())});
234+
235+
// DEFAULT accepts the first/last offsets; FULL rejects the intermediate decrease.
236+
ArrowArray root;
237+
ASSERT_EQ(ArrowArrayInitFromType(&root, NANOARROW_TYPE_STRUCT), NANOARROW_OK);
238+
ASSERT_EQ(ArrowArrayAllocateChildren(&root, 1), NANOARROW_OK);
239+
240+
ArrowArray* value_array = root.children[0];
241+
ASSERT_EQ(ArrowArrayInitFromType(value_array, NANOARROW_TYPE_STRING), NANOARROW_OK);
242+
const std::array<int32_t, 3> raw_offsets = {0, 5, 2};
243+
ASSERT_EQ(ArrowBufferAppend(ArrowArrayBuffer(value_array, 1), raw_offsets.data(),
244+
raw_offsets.size() * sizeof(int32_t)),
245+
NANOARROW_OK);
246+
ASSERT_EQ(ArrowBufferAppend(ArrowArrayBuffer(value_array, 2), "ab", 2),
247+
NANOARROW_OK);
248+
value_array->length = 2;
249+
value_array->null_count = 0;
250+
ASSERT_EQ(ArrowArrayFinishBuildingDefault(value_array, nullptr), NANOARROW_OK);
251+
252+
root.length = 2;
253+
root.null_count = 0;
254+
ASSERT_EQ(ArrowArrayFinishBuildingDefault(&root, nullptr), NANOARROW_OK);
255+
256+
auto projection =
257+
ProjectionContext::Make(*schema, *schema, function, validate_batch);
258+
ASSERT_THAT(projection, IsOk());
259+
260+
// Empty selection skips the content check without reading corrupt data.
261+
auto result = ProjectBatch(&root, /*row_indices=*/{}, projection.value());
262+
if (validate_batch) {
263+
EXPECT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
264+
EXPECT_THAT(result, HasErrorMessage("element size >= 0"));
265+
} else {
266+
EXPECT_THAT(result, IsOk());
267+
}
268+
if (result.has_value()) {
269+
ArrowArrayGuard guard(&result.value());
270+
EXPECT_EQ(result->length, 0);
271+
}
272+
EXPECT_EQ(root.release, nullptr);
273+
}
274+
}
275+
}
276+
277+
TEST(ProjectBatchTest, SkippedContentValidationStillRejectsInvalidStructure) {
278+
auto schema = MakeFullSchema();
279+
auto batch = MakeBatch(*schema, R"([[1,"a",1.0]])");
280+
for (auto function :
281+
{ProjectionContext::ProjectBatchFunction{nullptr}, ArrowComputeFunction()}) {
282+
ArrowArray array;
283+
ASSERT_TRUE(::arrow::ExportRecordBatch(*batch, &array).ok());
284+
array.length = -1;
285+
ICEBERG_UNWRAP_OR_FAIL(
286+
auto projection,
287+
ProjectionContext::Make(*schema, *schema, function, /*validate_batch=*/false));
288+
auto result = ProjectBatch(&array, {}, projection);
289+
EXPECT_THAT(result, IsError(ErrorKind::kInvalidArrowData));
290+
EXPECT_THAT(result, HasErrorMessage("length"));
291+
EXPECT_EQ(array.release, nullptr);
292+
}
293+
}
294+
295+
TEST(ProjectBatchTest, MovePreservesValidateBatch) {
296+
auto schema = MakeFullSchema();
297+
ICEBERG_UNWRAP_OR_FAIL(
298+
auto without_validation,
299+
ProjectionContext::Make(*schema, *schema, nullptr, /*validate_batch=*/false));
300+
auto moved = std::move(without_validation);
301+
EXPECT_FALSE(moved.validate_batch());
302+
ICEBERG_UNWRAP_OR_FAIL(auto validation,
303+
ProjectionContext::Make(*schema, *schema, nullptr));
304+
EXPECT_TRUE(validation.validate_batch());
305+
validation = std::move(moved);
306+
EXPECT_FALSE(validation.validate_batch());
307+
}
308+
234309
TEST(WriteProjectionTest, IdenticalSchemasTransferBatchWithoutConversion) {
235310
auto input_schema = MakeFullSchema();
236311
auto write_schema = MakeFullSchema();

‎src/iceberg/test/file_scan_task_reader_test.cc‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -671,6 +671,49 @@ TEST_F(FileScanTaskReaderTest, OpenWithDeletionVectorFiltersRows) {
671671
ASSERT_NO_FATAL_FAILURE(VerifyStream(&stream, R"([[1, "Foo"], [3, "Baz"]])"));
672672
}
673673

674+
TEST_F(FileScanTaskReaderTest, ReadsArrowValidateBatchProperty) {
675+
ICEBERG_UNWRAP_OR_FAIL(
676+
auto data_file,
677+
MakeDataFile(table_schema_,
678+
R"([[1, "Foo", "blue"], [2, "Bar", "red"], [3, "Baz", "green"]])"));
679+
ICEBERG_UNWRAP_OR_FAIL(
680+
auto pos_delete, MakePositionDeleteFile(CreateNewTempFilePathWithSuffix(".parquet"),
681+
{1}, data_file->file_path));
682+
FileScanTask task(data_file, {pos_delete});
683+
684+
struct RestoreProjection {
685+
ProjectionContext::ProjectBatchFunction function =
686+
ProjectionContext::ResolveProjectBatchFunction();
687+
~RestoreProjection() { ProjectionContext::RegisterProjectBatchFunction(function); }
688+
} restore;
689+
ProjectionContext::RegisterProjectBatchFunction(
690+
[](ArrowArray*, std::span<const int32_t>,
691+
ProjectionContext& projection) -> Result<ArrowArray> {
692+
return InvalidArgument("projection mode: {}",
693+
projection.validate_batch() ? "full" : "basic");
694+
});
695+
696+
for (bool validate_batch : {true, false}) {
697+
FileScanTaskReader::Options options{
698+
.io = file_io_,
699+
.table_schema = table_schema_,
700+
.projected_schema = projected_schema_,
701+
};
702+
if (!validate_batch) {
703+
options.properties[ReaderProperties::kArrowValidateBatch.key()] = "false";
704+
}
705+
ICEBERG_UNWRAP_OR_FAIL(auto reader, FileScanTaskReader::Make(std::move(options)));
706+
ICEBERG_UNWRAP_OR_FAIL(auto stream, reader->Open(task));
707+
708+
ArrowArray array{};
709+
EXPECT_NE(stream.get_next(&stream, &array), 0);
710+
EXPECT_THAT(stream.get_last_error(&stream),
711+
::testing::HasSubstr(validate_batch ? "projection mode: full"
712+
: "projection mode: basic"));
713+
stream.release(&stream);
714+
}
715+
}
716+
674717
TEST_F(FileScanTaskReaderTest, OpenWithEqualityDeletesAddsAndPrunesDeleteOnlyColumns) {
675718
ICEBERG_UNWRAP_OR_FAIL(
676719
auto data_file,

0 commit comments

Comments
 (0)