From 3d81f410d223669453343828c66d1df7243ae2d5 Mon Sep 17 00:00:00 2001 From: Qi Zhu <821684824@qq.com> Date: Tue, 28 Jul 2026 16:15:14 +0800 Subject: [PATCH] parquet: tolerate mixed row-group ordinal metadata on read Restore pre-#8715 tolerance for spec-valid parquet files whose row groups disagree on whether the optional RowGroup.ordinal field is populated (commonly produced by Go parquet writers flushing row groups incrementally). Decode no longer errors; instead: - all ordinals present: honored as written - none present: sequential-filled at decode time, unconditionally, so downstream consumers behave identically on fresh vs reused metadata - mixed: metadata left untouched; consumers that require complete ordinals fail deterministically instead of depending on which row groups a query selects RowNumberReader now rejects mixed-ordinal files up front (even when every selected row group carries an ordinal), and the encrypted column-metadata path returns a proper error instead of panicking on a missing ordinal. --- parquet/src/arrow/array_reader/row_number.rs | 80 +++++++- parquet/src/arrow/arrow_reader/mod.rs | 79 ++++++++ .../src/file/metadata/thrift/encryption.rs | 13 +- parquet/src/file/metadata/thrift/mod.rs | 183 ++++++++++++------ 4 files changed, 290 insertions(+), 65 deletions(-) diff --git a/parquet/src/arrow/array_reader/row_number.rs b/parquet/src/arrow/array_reader/row_number.rs index def88641a15a..dea1eab10e6c 100644 --- a/parquet/src/arrow/array_reader/row_number.rs +++ b/parquet/src/arrow/array_reader/row_number.rs @@ -47,14 +47,30 @@ impl RowNumberReader { // This is O(M) where M is the total number of row groups in the file let mut ordinal_to_offset: HashMap = HashMap::new(); let mut first_row_index: i64 = 0; + let mut missing_ordinals: usize = 0; for rg in parquet_metadata.row_groups() { if let Some(ordinal) = rg.ordinal() { ordinal_to_offset.insert(ordinal, first_row_index); + } else { + missing_ordinals += 1; } first_row_index += rg.num_rows(); } + // Mixed ordinals: refuse to compute row numbers for the whole file, + // even if every *selected* row group carries an ordinal. Otherwise the + // same file would yield row numbers or an error depending on which row + // groups a query's pruning happens to select. + if missing_ordinals > 0 && !ordinal_to_offset.is_empty() { + return Err(ParquetError::General(format!( + "Cannot compute row numbers: file has inconsistent row-group \ + ordinals ({} of {} row groups are missing the ordinal field)", + missing_ordinals, + parquet_metadata.num_row_groups(), + ))); + } + // Pass 2: Build ranges in the order specified by the row_groups iterator // This is O(N) where N is the number of selected row groups // This preserves the user's requested order instead of sorting by ordinal @@ -182,6 +198,15 @@ mod tests { } fn create_test_parquet_metadata(row_groups: Vec<(i16, i64)>) -> ParquetMetaData { + create_test_parquet_metadata_opt( + row_groups + .into_iter() + .map(|(ordinal, num_rows)| (Some(ordinal), num_rows)) + .collect(), + ) + } + + fn create_test_parquet_metadata_opt(row_groups: Vec<(Option, i64)>) -> ParquetMetaData { let schema_descr = create_test_schema(); let mut row_group_metas = vec![]; @@ -192,14 +217,14 @@ mod tests { .map(|col| ColumnChunkMetaData::builder(col.clone()).build().unwrap()) .collect(); - let row_group = RowGroupMetaData::builder(schema_descr.clone()) + let mut builder = RowGroupMetaData::builder(schema_descr.clone()) .set_num_rows(num_rows) - .set_ordinal(ordinal) .set_total_byte_size(100) - .set_column_metadata(columns) - .build() - .unwrap(); - row_group_metas.push(row_group); + .set_column_metadata(columns); + if let Some(ordinal) = ordinal { + builder = builder.set_ordinal(ordinal); + } + row_group_metas.push(builder.build().unwrap()); } let total_rows: i64 = row_group_metas.iter().map(|rg| rg.num_rows()).sum(); @@ -302,4 +327,47 @@ mod tests { assert_eq!(reader.read_records(4).unwrap(), 4); assert_eq!(consume_row_numbers(&mut reader), vec![3, 4, 5, 6]); } + + /// A file with *mixed* row-group ordinals must never produce row numbers — + /// even when every selected row group carries an ordinal. Otherwise the + /// same file would succeed or fail depending on which row groups a query's + /// pruning selects. + #[test] + fn test_mixed_ordinals_always_error() { + let metadata = create_test_parquet_metadata_opt(vec![ + (Some(0), 2), // has ordinal + (None, 2), // missing + (Some(2), 2), // has ordinal + ]); + + // Select only row groups WITH ordinals — must still fail. + let selected = vec![&metadata.row_groups()[0], &metadata.row_groups()[2]]; + let Err(err) = RowNumberReader::try_new(&metadata, selected.into_iter()) else { + panic!("mixed ordinals with ordinal-only selection must fail"); + }; + assert!( + err.to_string().contains("inconsistent row-group ordinals"), + "unexpected error: {err}" + ); + + // Selecting a row group without an ordinal fails too. + let selected = vec![&metadata.row_groups()[1]]; + assert!(RowNumberReader::try_new(&metadata, selected.into_iter()).is_err()); + } + + /// A file where *no* row group carries an ordinal fails for any selection + /// (fresh decode sequential-fills, so this only occurs for + /// programmatically-built metadata). + #[test] + fn test_no_ordinals_error() { + let metadata = create_test_parquet_metadata_opt(vec![(None, 2), (None, 2)]); + let selected = vec![&metadata.row_groups()[0]]; + let Err(err) = RowNumberReader::try_new(&metadata, selected.into_iter()) else { + panic!("row numbers without any ordinals must fail"); + }; + assert!( + err.to_string().contains("missing ordinal field"), + "unexpected error: {err}" + ); + } } diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index a9cd82b3a3de..ee1c4f8945d3 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/parquet/src/arrow/arrow_reader/mod.rs @@ -5614,6 +5614,85 @@ pub(crate) mod tests { Ok(()) } + /// A file with *mixed* row-group ordinal metadata (spec-valid — the + /// `RowGroup.ordinal` thrift field is optional; Go parquet writers emit + /// such files) must read fine without row numbers, and must fail + /// deterministically with them — even when every *selected* row group + /// carries an ordinal. See #10381. + #[test] + fn test_mixed_row_group_ordinals() -> Result<()> { + use crate::file::metadata::{ParquetMetaDataReader, RowGroupMetaData}; + + // 100 rows split across 4 row groups of 25 + let array = Int64Array::from_iter_values(5000..5100); + let batch = RecordBatch::try_from_iter([("col", Arc::new(array) as ArrayRef)])?; + let mut buffer = Vec::new(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(25)) + .build(); + let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema().clone(), Some(props))?; + for batch_chunk in (0..10).map(|i| batch.slice(i * 10, 10)) { + writer.write(&batch_chunk)?; + } + writer.close()?; + let buffer = Bytes::from(buffer); + + // Strip the ordinal from row group 1 to simulate a mixed-ordinal + // writer (the builder starts with no ordinal; copy everything else). + let metadata = ParquetMetaDataReader::new().parse_and_finish(&buffer)?; + let schema_descr = metadata.file_metadata().schema_descr_ptr(); + let mut row_groups = metadata.row_groups().to_vec(); + let stripped = row_groups[1].clone(); + let mut builder = RowGroupMetaData::builder(schema_descr) + .set_num_rows(stripped.num_rows()) + .set_total_byte_size(stripped.total_byte_size()) + .set_sorting_columns(stripped.sorting_columns().cloned()) + .set_column_metadata(stripped.columns().to_vec()); + if let Some(offset) = stripped.file_offset() { + builder = builder.set_file_offset(offset); + } + row_groups[1] = builder.build()?; + assert_eq!(row_groups[1].ordinal(), None); + let metadata = Arc::new(metadata.into_builder().set_row_groups(row_groups).build()); + + // Plain read (no row numbers): succeeds and returns all values. + let arrow_metadata = + ArrowReaderMetadata::try_new(Arc::clone(&metadata), ArrowReaderOptions::new())?; + let reader = + ParquetRecordBatchReaderBuilder::new_with_metadata(buffer.clone(), arrow_metadata) + .build()?; + let values: Vec = reader + .flat_map(|batch| { + let batch = batch.expect("could not read batch"); + batch + .column(0) + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + assert_eq!(values, (5000..5100).collect::>()); + + // Row-number read: fails deterministically, even when selecting only + // row groups that DO carry ordinals. + let row_number_field = Arc::new( + Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber), + ); + let options = ArrowReaderOptions::new().with_virtual_columns(vec![row_number_field])?; + let arrow_metadata = ArrowReaderMetadata::try_new(Arc::clone(&metadata), options)?; + let result = ParquetRecordBatchReaderBuilder::new_with_metadata(buffer, arrow_metadata) + .with_row_groups(vec![0]) // row group 0 has an ordinal + .build() + .and_then(|mut reader| reader.next().transpose().map_err(|e| e.into())); + let err = result.expect_err("row numbers over mixed ordinals must fail"); + assert!( + err.to_string().contains("inconsistent row-group ordinals"), + "unexpected error: {err}" + ); + + Ok(()) + } + #[derive(Debug, PartialEq)] struct ValuesAndRowNumbers { values: Vec, diff --git a/parquet/src/file/metadata/thrift/encryption.rs b/parquet/src/file/metadata/thrift/encryption.rs index 37e91ba99cba..1da40e8a064f 100644 --- a/parquet/src/file/metadata/thrift/encryption.rs +++ b/parquet/src/file/metadata/thrift/encryption.rs @@ -163,10 +163,21 @@ fn row_group_from_encrypted_thrift( } }; + // The ordinal is part of the AAD for encrypted column metadata. + // It can be missing here only for files with *mixed* row-group + // ordinals, which decode leaves untouched — fail cleanly rather + // than panic (see `ensure_row_group_ordinals`). + let rg_ordinal = rg.ordinal.ok_or_else(|| { + general_err!( + "Row group ordinal is required to decrypt column metadata for \ + column '{}', but the file's row-group ordinals are inconsistent", + d.path().string() + ) + })?; let column_aad = crate::encryption::modules::create_module_aad( decryptor.file_aad(), crate::encryption::modules::ModuleType::ColumnMetaData, - rg.ordinal.unwrap() as usize, + rg_ordinal as usize, i, None, )?; diff --git a/parquet/src/file/metadata/thrift/mod.rs b/parquet/src/file/metadata/thrift/mod.rs index d5a0112a5e1a..6331788666ea 100644 --- a/parquet/src/file/metadata/thrift/mod.rs +++ b/parquet/src/file/metadata/thrift/mod.rs @@ -814,17 +814,10 @@ pub(crate) fn parquet_metadata_from_bytes( validate_list_type(ElementType::Struct, &list_ident)?; let mut rg_vec = Vec::with_capacity(list_ident.size as usize); - // Read row groups and handle ordinal assignment - let mut assigner = OrdinalAssigner::new(); - for ordinal in 0..list_ident.size { - let ordinal: i16 = ordinal.try_into().map_err(|_| { - ParquetError::General(format!( - "Row group ordinal {ordinal} exceeds i16 max value", - )) - })?; - let rg = read_row_group(&mut prot, schema_descr, options)?; - rg_vec.push(assigner.ensure(ordinal, rg)?); + for _ in 0..list_ident.size { + rg_vec.push(read_row_group(&mut prot, schema_descr, options)?); } + ensure_row_group_ordinals(&mut rg_vec)?; row_groups = Some(rg_vec); } 5 => { @@ -915,56 +908,37 @@ pub(crate) fn parquet_metadata_from_bytes( Ok(ParquetMetaData::new(fmd, row_groups)) } -/// Assign [`RowGroupMetaData::ordinal`] if it is missing. -#[derive(Debug, Default)] -pub(crate) struct OrdinalAssigner { - first_has_ordinal: Option, -} - -impl OrdinalAssigner { - fn new() -> Self { - Default::default() +/// Ensure [`RowGroupMetaData::ordinal`] is usable after decode without +/// rejecting spec-valid files (`RowGroup.ordinal` is optional in the +/// parquet-format Thrift definition, with no uniformity requirement): +/// +/// - **All row groups carry ordinals** → honor them as written. +/// - **No row group carries an ordinal** → assign each row group its +/// position in the file. This happens unconditionally (not only when a +/// consumer needs it) so downstream users of the ordinal — the row +/// number virtual column, encryption chunk-key lookup — behave the same +/// whether the metadata was decoded fresh or reused from a prior read. +/// - **Mixed** → leave the metadata untouched. Positional backfill could +/// disagree with the ordinals that are present, and a partial backfill +/// would make row-number results depend on which row groups a query +/// happens to select. Consumers that require complete ordinals fail +/// deterministically instead (see `RowNumberReader::try_new`); plain +/// reads that never touch ordinals succeed. +fn ensure_row_group_ordinals(row_groups: &mut [RowGroupMetaData]) -> Result<()> { + if row_groups.iter().all(|rg| rg.ordinal.is_some()) { + return Ok(()); } - - /// Sets [`RowGroupMetaData::ordinal`] if it is missing. - /// - /// # Arguments - /// - actual_ordinal: The ordinal (index) of the row group being processed - /// in the file metadata. - /// - rg: The [`RowGroupMetaData`] to potentially modify. - /// - /// Ensures: - /// 1. If the first row group has an ordinal, all subsequent row groups must - /// also have ordinals. - /// 2. If the first row group does NOT have an ordinal, all subsequent row - /// groups must also not have ordinals. - fn ensure( - &mut self, - actual_ordinal: i16, - mut rg: RowGroupMetaData, - ) -> Result { - let rg_has_ordinal = rg.ordinal.is_some(); - - // Only set first_has_ordinal if it's None (first row group that arrives) - if self.first_has_ordinal.is_none() { - self.first_has_ordinal = Some(rg_has_ordinal); - } - - // assign ordinal if missing and consistent with first row group - let first_has_ordinal = self.first_has_ordinal.unwrap(); - if !first_has_ordinal && !rg_has_ordinal { - rg.ordinal = Some(actual_ordinal); - } else if first_has_ordinal != rg_has_ordinal { - return Err(general_err!( - "Inconsistent ordinal assignment: first_has_ordinal is set to \ - {} but row-group with actual ordinal {} has rg_has_ordinal set to {}", - first_has_ordinal, - actual_ordinal, - rg_has_ordinal - )); - } - Ok(rg) + if row_groups.iter().any(|rg| rg.ordinal.is_some()) { + // Mixed: leave as-is. + return Ok(()); } + for (idx, rg) in row_groups.iter_mut().enumerate() { + let ordinal: i16 = idx + .try_into() + .map_err(|_| general_err!("Row group ordinal {} exceeds i16 max value", idx))?; + rg.ordinal = Some(ordinal); + } + Ok(()) } thrift_struct!( @@ -2007,4 +1981,97 @@ pub(crate) mod tests { .expect_err("malformed bool field should return an error"); assert_malformed_bool_error(err); } + + /// Round-trip [`crate::file::metadata::ParquetMetaData`] with the given + /// per-row-group ordinals through thrift encode → decode, returning the + /// decoded ordinals. Exercises `ensure_row_group_ordinals`. + fn roundtrip_rg_ordinals(ordinals: &[Option]) -> Vec> { + use crate::file::metadata::ParquetMetaDataWriter; + use crate::file::metadata::{FileMetaData, ParquetMetaData, ParquetMetaDataReader}; + use crate::schema::types::Type as SchemaType; + + let field = SchemaType::primitive_type_builder("c", PhysicalType::INT32) + .build() + .unwrap(); + let schema = SchemaType::group_type_builder("schema") + .with_fields(vec![Arc::new(field)]) + .build() + .unwrap(); + let schema_descr = Arc::new(SchemaDescriptor::new(Arc::new(schema))); + + let row_groups = ordinals + .iter() + .map(|ordinal| { + let columns = schema_descr + .columns() + .iter() + .map(|col| ColumnChunkMetaData::builder(col.clone()).build().unwrap()) + .collect(); + let mut builder = + crate::file::metadata::RowGroupMetaData::builder(schema_descr.clone()) + .set_num_rows(10) + .set_total_byte_size(100) + .set_column_metadata(columns); + if let Some(ordinal) = ordinal { + builder = builder.set_ordinal(*ordinal); + } + builder.build().unwrap() + }) + .collect(); + + let file_metadata = FileMetaData::new( + 1, + 10 * ordinals.len() as i64, + None, + None, + schema_descr, + None, + ); + let metadata = ParquetMetaData::new(file_metadata, row_groups); + + let mut buffer = Vec::new(); + ParquetMetaDataWriter::new(&mut buffer, &metadata) + .finish() + .unwrap(); + // strip the 8-byte footer tail (length + magic) + let decoded = ParquetMetaDataReader::decode_metadata(&buffer[..buffer.len() - 8]).unwrap(); + decoded.row_groups().iter().map(|rg| rg.ordinal()).collect() + } + + /// All row groups carry ordinals: honored as written, even when they do + /// not match file position. + #[test] + fn ordinals_all_present_are_honored() { + assert_eq!( + roundtrip_rg_ordinals(&[Some(5), Some(1), Some(3)]), + vec![Some(5), Some(1), Some(3)], + ); + } + + /// No row group carries an ordinal: sequential-filled at decode time so + /// downstream consumers behave identically on fresh vs reused metadata. + #[test] + fn ordinals_none_present_are_sequentially_filled() { + assert_eq!( + roundtrip_rg_ordinals(&[None, None, None]), + vec![Some(0), Some(1), Some(2)], + ); + } + + /// Mixed ordinals (spec-valid; produced by e.g. Go parquet writers): + /// decode succeeds — the pre-#8715 behavior restored by #10381 — and the + /// metadata is left untouched so row numbering fails deterministically + /// rather than producing numbers that depend on row-group selection. + #[test] + fn ordinals_mixed_decode_succeeds_untouched() { + assert_eq!( + roundtrip_rg_ordinals(&[Some(0), None, Some(2)]), + vec![Some(0), None, Some(2)], + ); + // first missing, rest present — the exact Go-writer shape from #10381 + assert_eq!( + roundtrip_rg_ordinals(&[None, Some(1), Some(2)]), + vec![None, Some(1), Some(2)], + ); + } }