Skip to content

Commit 40f93ab

Browse files
ryzhykryzhyk
committed
[adapters] Keep container field ids when coercing INT96
`coerce_int96_to_resolution` rebuilds every struct, list and map field with `Field::new_struct`, `Field::new_list` and `Field::new`, which start from empty metadata. Leaf fields go through `field_with_new_type`, which clones, so a file holding any INT96 column comes back with the metadata stripped from its container fields alone. Reported upstream as apache/datafusion#24786. `PARQUET:field_id` lives in that metadata, and a `columnMapping` table's columns are paired with the read schema by id. Two ways that bites: - a table whose file names columns logically, a Uniform table written by an Iceberg writer, needs the id even at the top level: the column is dropped from the projection and then null-filled; - an ordinary Delta table names its file's columns `col-<id>`, so a top-level column still pairs by name, but `realign_array` pairs nested children by id alone, on purpose, since pairing them by position could match unrelated columns. A struct nested in a struct therefore arrives NULL. Either way it is silent data loss, or an error when the column is NOT NULL. Put the original metadata back on the coerced schema and cover both shapes. Without the restore each test null-fills its struct.
1 parent be3f074 commit 40f93ab

3 files changed

Lines changed: 350 additions & 6 deletions

File tree

‎crates/adapters/src/format/parquet.rs‎

Lines changed: 63 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -85,20 +85,78 @@ pub(crate) fn int96_as_micros(
8585
metadata: &ArrowReaderMetadata,
8686
options: &ArrowReaderOptions,
8787
) -> Result<Option<ArrowReaderMetadata>, ParquetError> {
88-
let Some(schema) = coerce_int96_to_resolution(
89-
metadata.parquet_schema(),
90-
metadata.schema(),
91-
&INT96_TIME_UNIT,
92-
) else {
88+
let original = metadata.schema();
89+
let Some(coerced) =
90+
coerce_int96_to_resolution(metadata.parquet_schema(), original, &INT96_TIME_UNIT)
91+
else {
9392
return Ok(None);
9493
};
94+
let schema = Schema::new_with_metadata(
95+
restore_field_metadata(original.fields(), coerced.fields()),
96+
coerced.metadata().clone(),
97+
);
9598
ArrowReaderMetadata::try_new(
9699
Arc::clone(metadata.metadata()),
97100
options.clone().with_schema(Arc::new(schema)),
98101
)
99102
.map(Some)
100103
}
101104

105+
/// Put back the field metadata that [`coerce_int96_to_resolution`] drops
106+
/// (<https://github.com/apache/datafusion/issues/24786>).
107+
///
108+
/// It rebuilds every struct, list and map field with `Field::new_struct`,
109+
/// `Field::new_list` and `Field::new`, which start from empty metadata, so any
110+
/// file holding an INT96 column comes back with its container fields stripped.
111+
/// Leaf fields keep theirs.
112+
///
113+
/// `PARQUET:field_id` lives in that metadata, and the Delta reader pairs a
114+
/// `columnMapping` table's columns with the read schema by id. A top-level
115+
/// column still survives on its name when the file uses the physical
116+
/// `col-<id>` names, but a nested one does not: `realign_array` pairs struct
117+
/// children by id alone, since pairing them by position could match unrelated
118+
/// columns. Either way the column turns into NULLs, or into an error when the
119+
/// table declares it NOT NULL.
120+
///
121+
/// Coercion changes types, never names or nesting, so the two field lists line
122+
/// up one to one.
123+
fn restore_field_metadata(original: &Fields, coerced: &Fields) -> Fields {
124+
original
125+
.iter()
126+
.zip(coerced.iter())
127+
.map(|(original, coerced)| {
128+
let data_type = match (original.data_type(), coerced.data_type()) {
129+
(DataType::Struct(original), DataType::Struct(coerced)) => {
130+
DataType::Struct(restore_field_metadata(original, coerced))
131+
}
132+
(DataType::List(original), DataType::List(coerced)) => DataType::List(
133+
restore_field_metadata(
134+
&[Arc::clone(original)].into(),
135+
&[Arc::clone(coerced)].into(),
136+
)[0]
137+
.clone(),
138+
),
139+
(DataType::Map(original, sorted), DataType::Map(coerced, _)) => DataType::Map(
140+
restore_field_metadata(
141+
&[Arc::clone(original)].into(),
142+
&[Arc::clone(coerced)].into(),
143+
)[0]
144+
.clone(),
145+
*sorted,
146+
),
147+
_ => coerced.data_type().clone(),
148+
};
149+
Arc::new(
150+
coerced
151+
.as_ref()
152+
.clone()
153+
.with_data_type(data_type)
154+
.with_metadata(original.metadata().clone()),
155+
)
156+
})
157+
.collect()
158+
}
159+
102160
/// Open `data` for reading.
103161
fn open_parquet_reader(
104162
data: Bytes,

‎crates/adapters/src/integrated/delta_table/test.rs‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4577,3 +4577,192 @@ fn delta_table_input_mixed_timestamp_encodings_test() {
45774577
timestamp_zset(&data)
45784578
);
45794579
}
4580+
4581+
/// A `columnMapping` table keeps its data files named logically and matches
4582+
/// them to the read schema, which names columns `col-<id>`, by Parquet field
4583+
/// id. An INT96 timestamp elsewhere in the file must not cost a struct column
4584+
/// its id, or the column comes through as NULLs.
4585+
#[tokio::test]
4586+
async fn delta_table_input_int96_preserves_column_mapping_ids_test() {
4587+
use crate::integrated::delta_table::deletion_vector::{ReadMode, filtered_parquet_table};
4588+
use crate::test::parquet_timestamps::write_column_mapped_int96_parquet;
4589+
use arrow::array::{Array, StringArray};
4590+
use arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField, TimeUnit};
4591+
use delta_kernel::object_store::local::LocalFileSystem;
4592+
use deltalake::{ObjectStore, Path as ObjectStorePath};
4593+
use roaring::RoaringTreemap;
4594+
4595+
fn mapped(field: ArrowField, id: &str) -> ArrowField {
4596+
field.with_metadata(HashMap::from([(
4597+
"delta.columnMapping.id".to_string(),
4598+
id.to_string(),
4599+
)]))
4600+
}
4601+
4602+
let table_dir = TempDir::new().unwrap();
4603+
let file_name = "data.parquet";
4604+
let ts = timestamp_test_data()[1].ts.unwrap();
4605+
write_column_mapped_int96_parquet(&table_dir.path().join(file_name), "hello", ts);
4606+
4607+
// The read schema of a column-mapped table: physical names, ids in metadata.
4608+
let logical_schema = Arc::new(ArrowSchema::new(vec![
4609+
mapped(ArrowField::new("col-1", ArrowDataType::Int64, false), "1"),
4610+
mapped(
4611+
ArrowField::new_struct(
4612+
"col-2",
4613+
vec![mapped(
4614+
ArrowField::new("col-3", ArrowDataType::Utf8, true),
4615+
"3",
4616+
)],
4617+
true,
4618+
),
4619+
"2",
4620+
),
4621+
mapped(
4622+
ArrowField::new(
4623+
"col-4",
4624+
ArrowDataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
4625+
true,
4626+
),
4627+
"4",
4628+
),
4629+
]));
4630+
4631+
let store: Arc<dyn ObjectStore> =
4632+
Arc::new(LocalFileSystem::new_with_prefix(table_dir.path()).unwrap());
4633+
let provider = filtered_parquet_table(
4634+
store,
4635+
ObjectStorePath::from(file_name),
4636+
RoaringTreemap::new(),
4637+
logical_schema,
4638+
ReadMode::NotInBitmap,
4639+
)
4640+
.await
4641+
.unwrap();
4642+
4643+
let batches = SessionContext::new()
4644+
.read_table(provider)
4645+
.unwrap()
4646+
.collect()
4647+
.await
4648+
.unwrap();
4649+
let batch = &batches[0];
4650+
4651+
let struct_column = batch
4652+
.column_by_name("col-2")
4653+
.unwrap()
4654+
.as_any()
4655+
.downcast_ref::<arrow::array::StructArray>()
4656+
.unwrap();
4657+
let label = struct_column
4658+
.column_by_name("col-3")
4659+
.unwrap()
4660+
.as_any()
4661+
.downcast_ref::<StringArray>()
4662+
.unwrap();
4663+
assert!(
4664+
!label.is_null(0),
4665+
"the struct column lost its field id and was null-filled"
4666+
);
4667+
assert_eq!(label.value(0), "hello");
4668+
4669+
// The timestamp that triggered the coercion still arrives intact.
4670+
let timestamps = batch
4671+
.column_by_name("col-4")
4672+
.unwrap()
4673+
.as_any()
4674+
.downcast_ref::<arrow::array::TimestampMicrosecondArray>()
4675+
.unwrap();
4676+
assert_eq!(timestamps.value(0), ts.microseconds());
4677+
}
4678+
4679+
/// A native Delta `columnMapping` table names its data file's columns by their
4680+
/// physical `col-<id>` name, so the reader resolves a top-level column by name
4681+
/// even without its field id. Nested children it resolves by id alone,
4682+
/// deliberately: matching them by position could pair unrelated columns. An
4683+
/// INT96 timestamp elsewhere in the file must therefore not cost a nested
4684+
/// struct its id, or that struct arrives NULL.
4685+
#[tokio::test]
4686+
async fn delta_table_input_int96_preserves_nested_column_mapping_ids_test() {
4687+
use crate::integrated::delta_table::deletion_vector::{ReadMode, filtered_parquet_table};
4688+
use crate::test::parquet_timestamps::write_nested_column_mapped_int96_parquet;
4689+
use arrow::array::{Array, StringArray, StructArray};
4690+
use arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField};
4691+
use delta_kernel::object_store::local::LocalFileSystem;
4692+
use deltalake::{ObjectStore, Path as ObjectStorePath};
4693+
use roaring::RoaringTreemap;
4694+
4695+
fn mapped(field: ArrowField, id: &str) -> ArrowField {
4696+
field.with_metadata(HashMap::from([(
4697+
"delta.columnMapping.id".to_string(),
4698+
id.to_string(),
4699+
)]))
4700+
}
4701+
4702+
let table_dir = TempDir::new().unwrap();
4703+
let file_name = "data.parquet";
4704+
let ts = timestamp_test_data()[1].ts.unwrap();
4705+
write_nested_column_mapped_int96_parquet(&table_dir.path().join(file_name), "deep", ts);
4706+
4707+
let logical_schema = Arc::new(ArrowSchema::new(vec![
4708+
mapped(ArrowField::new("col-1", ArrowDataType::Int64, false), "1"),
4709+
mapped(
4710+
ArrowField::new_struct(
4711+
"col-2",
4712+
vec![mapped(
4713+
ArrowField::new_struct(
4714+
"col-3",
4715+
vec![mapped(
4716+
ArrowField::new("col-4", ArrowDataType::Utf8, true),
4717+
"4",
4718+
)],
4719+
true,
4720+
),
4721+
"3",
4722+
)],
4723+
true,
4724+
),
4725+
"2",
4726+
),
4727+
]));
4728+
4729+
let store: Arc<dyn ObjectStore> =
4730+
Arc::new(LocalFileSystem::new_with_prefix(table_dir.path()).unwrap());
4731+
let provider = filtered_parquet_table(
4732+
store,
4733+
ObjectStorePath::from(file_name),
4734+
RoaringTreemap::new(),
4735+
logical_schema,
4736+
ReadMode::NotInBitmap,
4737+
)
4738+
.await
4739+
.unwrap();
4740+
4741+
let batches = SessionContext::new()
4742+
.read_table(provider)
4743+
.unwrap()
4744+
.collect()
4745+
.await
4746+
.unwrap();
4747+
let outer = batches[0]
4748+
.column_by_name("col-2")
4749+
.unwrap()
4750+
.as_any()
4751+
.downcast_ref::<StructArray>()
4752+
.unwrap();
4753+
let inner = outer.column_by_name("col-3").unwrap();
4754+
assert!(
4755+
!inner.is_null(0),
4756+
"the nested struct lost its field id and was null-filled"
4757+
);
4758+
let label = inner
4759+
.as_any()
4760+
.downcast_ref::<StructArray>()
4761+
.unwrap()
4762+
.column_by_name("col-4")
4763+
.unwrap()
4764+
.as_any()
4765+
.downcast_ref::<StringArray>()
4766+
.unwrap();
4767+
assert_eq!(label.value(0), "deep");
4768+
}

‎crates/adapters/src/test/parquet_timestamps.rs‎

Lines changed: 98 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ use std::path::{Path, PathBuf};
1717
use std::sync::Arc;
1818
use std::time::{SystemTime, UNIX_EPOCH};
1919

20-
use parquet::data_type::{Int64Type, Int96, Int96Type};
20+
use parquet::data_type::{ByteArray, ByteArrayType, Int64Type, Int96, Int96Type};
2121
use parquet::file::properties::WriterProperties;
2222
use parquet::file::writer::SerializedFileWriter;
2323
use parquet::schema::parser::parse_message_type;
@@ -286,6 +286,103 @@ pub fn write_nested_int96_parquet(path: &Path, rows: &[TimestampTestStruct]) {
286286
writer.close().unwrap();
287287
}
288288

289+
/// Write a file shaped like a `columnMapping` table's data file: columns named
290+
/// logically and stamped with Parquet field ids, next to an INT96 timestamp.
291+
///
292+
/// The reader matches such a file's columns to the Delta read schema, which
293+
/// names them `col-<id>`, by id alone.
294+
pub fn write_column_mapped_int96_parquet(path: &Path, label: &str, ts: Timestamp) {
295+
const MESSAGE_TYPE: &str = "
296+
message spark_schema {
297+
REQUIRED INT64 row_id = 1;
298+
OPTIONAL group info = 2 {
299+
OPTIONAL BYTE_ARRAY label (STRING) = 3;
300+
}
301+
OPTIONAL INT96 ts = 4;
302+
}
303+
";
304+
305+
let schema = Arc::new(parse_message_type(MESSAGE_TYPE).unwrap());
306+
let props = Arc::new(WriterProperties::builder().build());
307+
let mut writer = SerializedFileWriter::new(File::create(path).unwrap(), schema, props).unwrap();
308+
let mut row_group = writer.next_row_group().unwrap();
309+
310+
let mut column = row_group.next_column().unwrap().unwrap();
311+
column
312+
.typed::<Int64Type>()
313+
.write_batch(&[1], None, None)
314+
.unwrap();
315+
column.close().unwrap();
316+
317+
// `info` and `info.label` are both present, so definition level 2.
318+
let mut column = row_group.next_column().unwrap().unwrap();
319+
column
320+
.typed::<ByteArrayType>()
321+
.write_batch(&[ByteArray::from(label)], Some(&[2]), None)
322+
.unwrap();
323+
column.close().unwrap();
324+
325+
let mut column = row_group.next_column().unwrap().unwrap();
326+
column
327+
.typed::<Int96Type>()
328+
.write_batch(&[to_int96(ts, 0)], Some(&[1]), None)
329+
.unwrap();
330+
column.close().unwrap();
331+
332+
row_group.close().unwrap();
333+
writer.close().unwrap();
334+
}
335+
336+
/// Write a file shaped like a native Delta `columnMapping` table's: columns
337+
/// named by their physical `col-<id>` name, stamped with Parquet field ids, a
338+
/// struct nested inside a struct, and an INT96 timestamp.
339+
///
340+
/// The names already match the read schema here, so the reader resolves the
341+
/// outer column by name. Its nested children it resolves by id alone.
342+
pub fn write_nested_column_mapped_int96_parquet(path: &Path, label: &str, ts: Timestamp) {
343+
const MESSAGE_TYPE: &str = r#"
344+
message spark_schema {
345+
REQUIRED INT64 "col-1" = 1;
346+
OPTIONAL group "col-2" = 2 {
347+
OPTIONAL group "col-3" = 3 {
348+
OPTIONAL BYTE_ARRAY "col-4" (STRING) = 4;
349+
}
350+
}
351+
OPTIONAL INT96 "col-5" = 5;
352+
}
353+
"#;
354+
355+
let schema = Arc::new(parse_message_type(MESSAGE_TYPE).unwrap());
356+
let props = Arc::new(WriterProperties::builder().build());
357+
let mut writer = SerializedFileWriter::new(File::create(path).unwrap(), schema, props).unwrap();
358+
let mut row_group = writer.next_row_group().unwrap();
359+
360+
let mut column = row_group.next_column().unwrap().unwrap();
361+
column
362+
.typed::<Int64Type>()
363+
.write_batch(&[1], None, None)
364+
.unwrap();
365+
column.close().unwrap();
366+
367+
// `col-2`, `col-3` and `col-4` are all present, so definition level 3.
368+
let mut column = row_group.next_column().unwrap().unwrap();
369+
column
370+
.typed::<ByteArrayType>()
371+
.write_batch(&[ByteArray::from(label)], Some(&[3]), None)
372+
.unwrap();
373+
column.close().unwrap();
374+
375+
let mut column = row_group.next_column().unwrap().unwrap();
376+
column
377+
.typed::<Int96Type>()
378+
.write_batch(&[to_int96(ts, 0)], Some(&[1]), None)
379+
.unwrap();
380+
column.close().unwrap();
381+
382+
row_group.close().unwrap();
383+
writer.close().unwrap();
384+
}
385+
289386
/// Write one data file into `table_dir` and return its `add` action.
290387
fn add_data_file(
291388
table_dir: &Path,

0 commit comments

Comments
 (0)