From bba363a02315c1e0c3a58296b87b9edd2fee153f Mon Sep 17 00:00:00 2001 From: jiaqizho Date: Mon, 16 Mar 2026 19:36:09 +0800 Subject: [PATCH] fix: validate schema type compatibility when reading parquet files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When reading parquet files, we weren't checking whether the caller's read schema has compatible field types with the actual file schema. This could lead to subtle issues — for example, if someone passes a struct type for a field that's actually int64 in the file, Arrow's C Data Interface would happily import it with the wrong memory layout, which can cause memory corruption or ASAN violations. This adds a ValidateSchemaCompatibility check in ParquetFormatReader::open() that compares each field in the read schema against the file schema. If a field exists in both but has a different type, we now fail early with a clear error message instead of silently misinterpreting the data. The read_schema is threaded through from FormatReader::create() down to ParquetFormatReader so the validation happens right after we read the file's actual schema. Fields that only exist in the read schema (schema evolution case) are skipped. Signed-off-by: jiaqizho --- .../milvus-storage/common/arrow_util.h | 6 ++ .../format/parquet/parquet_format_reader.h | 2 + cpp/src/common/arrow_util.cpp | 20 ++++ cpp/src/format/format_reader.cpp | 2 +- .../format/parquet/parquet_format_reader.cpp | 13 ++- cpp/test/api_writer_reader_test.cpp | 92 +++++++++++++++++++ cpp/test/format/parquet/file_writer_test.cpp | 6 +- 7 files changed, 135 insertions(+), 6 deletions(-) diff --git a/cpp/include/milvus-storage/common/arrow_util.h b/cpp/include/milvus-storage/common/arrow_util.h index ee70b0658..35d47a063 100644 --- a/cpp/include/milvus-storage/common/arrow_util.h +++ b/cpp/include/milvus-storage/common/arrow_util.h @@ -65,4 +65,10 @@ arrow::Result GetEnvVar(const char* name); arrow::Result GetEnvVar(const std::string& name); +// Validate that the read schema's field types are compatible with the file schema. +// For each field in read_schema, if a field with the same name exists in file_schema, +// their types must match. Returns an error if any type mismatch is found. +arrow::Status ValidateSchemaCompatibility(const std::shared_ptr& read_schema, + const std::shared_ptr& file_schema); + } // namespace milvus_storage diff --git a/cpp/include/milvus-storage/format/parquet/parquet_format_reader.h b/cpp/include/milvus-storage/format/parquet/parquet_format_reader.h index be75a1115..e746b2ff6 100644 --- a/cpp/include/milvus-storage/format/parquet/parquet_format_reader.h +++ b/cpp/include/milvus-storage/format/parquet/parquet_format_reader.h @@ -34,6 +34,7 @@ class ParquetFormatReader final : public FormatReader { const milvus_storage::api::Properties& properties, const std::vector& needed_columns, const std::function& key_retriever, + const std::shared_ptr& read_schema = nullptr, uint64_t file_size = 0, uint64_t footer_size = 0); @@ -75,6 +76,7 @@ class ParquetFormatReader final : public FormatReader { std::string path_; std::shared_ptr fs_; std::shared_ptr schema_; + std::shared_ptr read_schema_; // optional: used for type validation against file schema milvus_storage::api::Properties properties_; std::vector needed_columns_; std::function key_retriever_; diff --git a/cpp/src/common/arrow_util.cpp b/cpp/src/common/arrow_util.cpp index e773057f3..c8932d7f1 100644 --- a/cpp/src/common/arrow_util.cpp +++ b/cpp/src/common/arrow_util.cpp @@ -204,4 +204,24 @@ arrow::Result GetEnvVar(const char* name) { arrow::Result GetEnvVar(const std::string& name) { return GetEnvVar(name.c_str()); } +arrow::Status ValidateSchemaCompatibility(const std::shared_ptr& read_schema, + const std::shared_ptr& file_schema) { + if (!read_schema || !file_schema) { + return arrow::Status::Invalid("read_schema or file_schema is null"); + } + for (int i = 0; i < read_schema->num_fields(); ++i) { + auto& read_field = read_schema->field(i); + auto file_field = file_schema->GetFieldByName(read_field->name()); + if (file_field == nullptr) { + continue; // field not in file, skip (schema evolution: new column) + } + if (!read_field->type()->Equals(file_field->type())) { + return arrow::Status::Invalid( + fmt::format("Schema type mismatch for field '{}': read schema has type '{}' but file has type '{}'", + read_field->name(), read_field->type()->ToString(), file_field->type()->ToString())); + } + } + return arrow::Status::OK(); +} + } // namespace milvus_storage diff --git a/cpp/src/format/format_reader.cpp b/cpp/src/format/format_reader.cpp index 7ca8751aa..449fe5dc7 100644 --- a/cpp/src/format/format_reader.cpp +++ b/cpp/src/format/format_reader.cpp @@ -45,7 +45,7 @@ arrow::Result> FormatReader::create( ARROW_ASSIGN_OR_RAISE(auto uri, StorageUri::Parse(file.path)); std::string resolved_path = uri.scheme.empty() ? file.path : uri.key; format_reader = std::make_shared( - file_system, resolved_path, properties, needed_columns, key_retriever, + file_system, resolved_path, properties, needed_columns, key_retriever, read_schema, file.Get(api::kPropertyFileSize), file.Get(api::kPropertyFooterSize)); } else if (format == LOON_FORMAT_VORTEX) { ARROW_ASSIGN_OR_RAISE(auto file_system, FilesystemCache::getInstance().get(properties, file.path)); diff --git a/cpp/src/format/parquet/parquet_format_reader.cpp b/cpp/src/format/parquet/parquet_format_reader.cpp index 00647f302..b37f3d684 100644 --- a/cpp/src/format/parquet/parquet_format_reader.cpp +++ b/cpp/src/format/parquet/parquet_format_reader.cpp @@ -99,13 +99,15 @@ ParquetFormatReader::ParquetFormatReader(const std::shared_ptr& needed_columns, const std::function& key_retriever, + const std::shared_ptr& read_schema, uint64_t file_size, uint64_t footer_size) : path_(path), fs_(fs), schema_(nullptr), - properties_(properties), - needed_columns_(needed_columns), + read_schema_(std::move(read_schema)), + properties_(std::move(properties)), + needed_columns_(std::move(needed_columns)), key_retriever_(key_retriever), file_size_(file_size), footer_size_(footer_size), @@ -214,6 +216,12 @@ arrow::Status ParquetFormatReader::open() { // get the schema and create needed column indices std::shared_ptr file_schema; ARROW_RETURN_NOT_OK(file_reader_->GetSchema(&file_schema)); + + // validate read schema type compatibility with file schema + if (read_schema_) { + ARROW_RETURN_NOT_OK(ValidateSchemaCompatibility(read_schema_, file_schema)); + } + schema_ = file_schema; // Convert needed column names to column indices @@ -495,6 +503,7 @@ ParquetFormatReader::ParquetFormatReader(const ParquetFormatReader& other, : path_(other.path_), fs_(other.fs_), schema_(other.schema_), + read_schema_(other.read_schema_), properties_(other.properties_), needed_columns_(other.needed_columns_), key_retriever_(other.key_retriever_), diff --git a/cpp/test/api_writer_reader_test.cpp b/cpp/test/api_writer_reader_test.cpp index 7c5b7731b..81b9e35b6 100644 --- a/cpp/test/api_writer_reader_test.cpp +++ b/cpp/test/api_writer_reader_test.cpp @@ -22,6 +22,8 @@ #include #include +#include +#include #include #include #include @@ -1505,6 +1507,96 @@ TEST_P(APIWriterReaderTest, ReadWithoutSchema) { } } +TEST_P(APIWriterReaderTest, SchemaTypeMismatchIntToStruct) { + // Write with schema {int64, int64} + auto write_schema = arrow::schema({ + arrow::field("field_a", arrow::int64(), false), + arrow::field("field_b", arrow::int64(), false), + }); + + // Build data + arrow::Int64Builder a_builder, b_builder; + for (int64_t i = 0; i < 10; ++i) { + ASSERT_STATUS_OK(a_builder.Append(i)); + ASSERT_STATUS_OK(b_builder.Append(i * 100)); + } + std::shared_ptr a_array, b_array; + ASSERT_STATUS_OK(a_builder.Finish(&a_array)); + ASSERT_STATUS_OK(b_builder.Finish(&b_array)); + auto write_batch = arrow::RecordBatch::Make(write_schema, 10, {a_array, b_array}); + + // Write + ASSERT_AND_ASSIGN(auto policy, CreateSinglePolicy(format, write_schema)); + auto writer = Writer::create(base_path_, write_schema, std::move(policy), properties_); + ASSERT_NE(writer, nullptr); + ASSERT_OK(writer->write(write_batch)); + ASSERT_AND_ASSIGN(auto cgs, writer->close()); + + // Read with schema {int64, struct{int64}} — type mismatch on field_b + auto read_schema = arrow::schema({ + arrow::field("field_a", arrow::int64(), false), + arrow::field("field_b", arrow::struct_({arrow::field("sub", arrow::int64())}), false), + }); + + auto reader = Reader::create(cgs, read_schema, nullptr, properties_); + ASSERT_NE(reader, nullptr); + + auto batch_reader_result = reader->get_record_batch_reader(); + if (!batch_reader_result.ok()) { + // Parquet: validation fails at open() time + return; + } + + // Vortex: validation fails at ReadNext() time + auto batch_reader = batch_reader_result.MoveValueUnsafe(); + std::shared_ptr batch; + auto status = batch_reader->ReadNext(&batch); + ASSERT_FALSE(status.ok()) << "Expected failure due to schema type mismatch (int64 vs struct)"; +} + +// Pure C Data Interface test: export int64 data, import with struct schema +TEST_P(APIWriterReaderTest, ImportWithWrongSchema) { + // Build a RecordBatch with {int64, int64} + auto real_schema = arrow::schema({ + arrow::field("field_a", arrow::int64(), false), + arrow::field("field_b", arrow::int64(), false), + }); + arrow::Int64Builder a_builder, b_builder; + for (int64_t i = 0; i < 10; ++i) { + ASSERT_STATUS_OK(a_builder.Append(i)); + ASSERT_STATUS_OK(b_builder.Append(i * 100)); + } + std::shared_ptr a_array, b_array; + ASSERT_STATUS_OK(a_builder.Finish(&a_array)); + ASSERT_STATUS_OK(b_builder.Finish(&b_array)); + auto batch = arrow::RecordBatch::Make(real_schema, 10, {a_array, b_array}); + + // Export only the ArrowArray (no schema) + struct ArrowArray c_array; + ASSERT_OK(arrow::ExportRecordBatch(*batch, &c_array)); + + // Import with wrong schema: field_b is struct instead of int64 + auto wrong_schema = arrow::schema({ + arrow::field("field_a", arrow::int64(), false), + arrow::field("field_b", arrow::struct_({arrow::field("sub", arrow::int64())}), false), + }); + + auto import_result = arrow::ImportRecordBatch(&c_array, wrong_schema); + if (!import_result.ok()) { + // Arrow caught the mismatch at import time - clean up the exported array + if (c_array.release != nullptr) { + c_array.release(&c_array); + } + SUCCEED() << "Import correctly rejected schema mismatch: " << import_result.status().ToString(); + return; + } + + // Arrow allowed the import despite the type mismatch - verify the data is at least accessible + auto imported = import_result.ValueOrDie(); + ASSERT_NE(imported, nullptr); + ASSERT_EQ(imported->num_columns(), 2); +} + INSTANTIATE_TEST_SUITE_P(APIWriterReaderTestP, APIWriterReaderTest, ::testing::Combine(::testing::Values(LOON_FORMAT_PARQUET, LOON_FORMAT_VORTEX), diff --git a/cpp/test/format/parquet/file_writer_test.cpp b/cpp/test/format/parquet/file_writer_test.cpp index db3454e42..75f6f81a1 100644 --- a/cpp/test/format/parquet/file_writer_test.cpp +++ b/cpp/test/format/parquet/file_writer_test.cpp @@ -556,9 +556,9 @@ TEST_F(ParquetFileWriterTest, FooterSizeNotMatch) { // The reader uses footer_size to pre-read the footer in a single IO; // if the size is wrong, it falls back to Arrow's normal 2-step footer read. auto verify_read = [&](uint64_t footer_size) { - auto reader = - milvus_storage::parquet::ParquetFormatReader(fs_, temp_file, properties_, /*needed_columns=*/{}, - /*key_retriever=*/nullptr, cached_file_size, footer_size); + auto reader = milvus_storage::parquet::ParquetFormatReader(fs_, temp_file, properties_, /*needed_columns=*/{}, + /*key_retriever=*/nullptr, /*read_schema=*/nullptr, + cached_file_size, footer_size); ASSERT_STATUS_OK(reader.open()); ASSERT_AND_ASSIGN(auto row_group_infos, reader.get_row_group_infos());