Skip to content

Commit 9959632

Browse files
ryzhykryzhyk
committed
[adapters] Read INT96 Parquet timestamps as microseconds
Spark and Databricks still write Delta timestamps as the deprecated INT96 type, which holds a Julian day plus nanoseconds within the day and so spans a wider date range than a 64-bit nanosecond count. arrow-rs decodes INT96 as nanoseconds by default, and `Int96::to_nanos` (in the parquet crate) wraps on overflow, so any date outside 1677-09-21 .. 2262-04-11 was ingested corrupted. Datafusion has a setting to convert INT96 timestamps to micros. Three reader paths need the setting separately: - the snapshot scan, via the connector's datafusion session, whose parquet options delta-rs copies into its own; - the per-file reads follow and CDC mode perform, whose `ParquetFormat` ignores the session settings and must be handed them; - the deletion-vector reader and the `parquet` input format, which drive arrow-rs directly. The follow and CDC reads take the session's whole Parquet configuration rather than the INT96 setting alone, so the two paths cannot disagree about how to decode a file. Delta pins the logical `timestamp` type at microsecond precision but leaves the physical encoding to the writer, so one table's files can mix INT96 with INT64 microseconds. Coercion is decided per file from that file's own Parquet schema, so INT64 files are left untouched. Coercing a file rebuilds its struct, list and map fields, and only datafusion 55.2 and later keep their Parquet field ids. Without them, a column-mapped table reads a struct column as NULL (apache/datafusion#24786). The connector's data type mapping in the docs claimed timestamp values are rounded to the nearest millisecond. They are kept to microsecond precision, so say that instead. Signed-off-by: Leonid Ryzhyk <[email protected]>
1 parent 9459f7c commit 9959632

9 files changed

Lines changed: 944 additions & 18 deletions

File tree

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

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use arrow::datatypes::{
88
TimeUnit,
99
};
1010
use bytes::Bytes;
11+
use datafusion::datasource::file_format::parquet::Int96Coercer;
1112
use dbsp::operator::StagedBuffers;
1213
use erased_serde::Serialize as ErasedSerialize;
1314
use feldera_adapterlib::ConnectorMetadata;
@@ -19,7 +20,11 @@ use feldera_types::serde_with_context::serde_config::{
1920
};
2021
use feldera_types::serde_with_context::{DateFormat, SqlSerdeConfig, TimeFormat, TimestampFormat};
2122
use parquet::arrow::ArrowWriter;
22-
use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
23+
use parquet::arrow::arrow_reader::{
24+
ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader,
25+
ParquetRecordBatchReaderBuilder,
26+
};
27+
use parquet::errors::ParquetError;
2328
use parquet::file::properties::WriterProperties;
2429
use serde::Deserialize;
2530
use serde_arrow::ArrayBuilder;
@@ -56,6 +61,59 @@ pub const fn default_arrow_serde_config() -> &'static SqlSerdeConfig {
5661
}
5762
}
5863

64+
/// Time unit Feldera decodes deprecated INT96 Parquet timestamps at.
65+
///
66+
/// Spark and Databricks still write timestamps as INT96, which spans a wider
67+
/// date range than 64-bit nanoseconds. arrow-rs converts INT96 to nanoseconds
68+
/// unless told otherwise, wrapping silently: 4000-12-31 becomes 2247-05-04.
69+
/// Microseconds cover the whole INT96 range and match Feldera's own TIMESTAMP
70+
/// precision, so avoiding the nanosecond conversion preserves the value.
71+
pub(crate) const INT96_TIME_UNIT: TimeUnit = TimeUnit::Microsecond;
72+
73+
/// [`INT96_TIME_UNIT`] as datafusion spells it in its
74+
/// `execution.parquet.coerce_int96` setting.
75+
pub(crate) const INT96_TIME_UNIT_NAME: &str = "us";
76+
77+
/// Retype `metadata`'s INT96 timestamp columns as [`INT96_TIME_UNIT`], or
78+
/// return `None` when the file declares no such column.
79+
///
80+
/// Rebuilding the metadata restates the options it was loaded with, so the
81+
/// caller passes the same `options` it used; defaulting them here would drop
82+
/// the caller's settings on exactly the files that contain an INT96 column.
83+
pub(crate) fn int96_as_micros(
84+
metadata: &ArrowReaderMetadata,
85+
options: &ArrowReaderOptions,
86+
) -> Result<Option<ArrowReaderMetadata>, ParquetError> {
87+
let Some(schema) = Int96Coercer::new(
88+
metadata.parquet_schema(),
89+
metadata.schema(),
90+
&INT96_TIME_UNIT,
91+
)
92+
.coerce() else {
93+
return Ok(None);
94+
};
95+
ArrowReaderMetadata::try_new(
96+
Arc::clone(metadata.metadata()),
97+
options.clone().with_schema(Arc::new(schema)),
98+
)
99+
.map(Some)
100+
}
101+
102+
/// Open `data` for reading.
103+
fn open_parquet_reader(
104+
data: Bytes,
105+
batch_size: usize,
106+
) -> Result<ParquetRecordBatchReader, ParquetError> {
107+
let options = ArrowReaderOptions::default();
108+
let mut metadata = ArrowReaderMetadata::load(&data, options.clone())?;
109+
if let Some(coerced) = int96_as_micros(&metadata, &options)? {
110+
metadata = coerced;
111+
}
112+
ParquetRecordBatchReaderBuilder::new_with_metadata(data, metadata)
113+
.with_batch_size(batch_size)
114+
.build()
115+
}
116+
59117
/// CSV format parser.
60118
pub struct ParquetInputFormat;
61119

@@ -114,7 +172,7 @@ impl Parser for ParquetParser {
114172

115173
let bytes = Bytes::copy_from_slice(data);
116174

117-
let parquet_reader = match ParquetRecordBatchReader::try_new(bytes, 1_000_000) {
175+
let parquet_reader = match open_parquet_reader(bytes, 1_000_000) {
118176
Ok(parquet_reader) => parquet_reader,
119177
Err(e) => {
120178
return (

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

Lines changed: 215 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,16 @@ use std::sync::{Arc, Mutex};
66
use std::thread::sleep;
77
use std::time::Duration;
88

9-
use arrow::array::RecordBatch;
10-
use arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema};
9+
use arrow::array::{Array, ListArray, RecordBatch, StructArray, TimestampMicrosecondArray};
10+
use arrow::datatypes::{
11+
DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, TimeUnit,
12+
};
13+
use bytes::Bytes;
14+
use chrono::DateTime;
1115
use dbsp::OrdZSet;
1216
use dbsp::utils::Tup2;
1317
use feldera_adapterlib::transport::OutputBatchType;
18+
use feldera_sqllib::Timestamp;
1419
use feldera_sqllib::Variant;
1520
use feldera_types::deserialize_table_record;
1621
use feldera_types::format::parquet::ParquetEncoderConfig;
@@ -30,9 +35,18 @@ use tracing_subscriber::util::SubscriberInitExt;
3035

3136
use crate::{
3237
catalog::SerBatchReader,
33-
format::{Encoder, parquet::ParquetEncoder},
38+
format::{Encoder, parquet::ParquetEncoder, parquet::open_parquet_reader},
3439
static_compile::seroutput::SerBatchImpl,
35-
test::{DEFAULT_TIMEOUT_MS, MockOutputConsumer, TestStruct2, mock_input_pipeline, wait},
40+
test::{
41+
DEFAULT_TIMEOUT_MS, MockOutputConsumer, TestStruct2, TimestampTestStruct,
42+
mock_input_pipeline,
43+
parquet_timestamps::{
44+
TimestampEncoding, timestamp_test_data, write_mixed_column_parquet,
45+
write_nested_int96_parquet, write_sub_microsecond_int96_parquet,
46+
write_timestamp_parquet,
47+
},
48+
wait,
49+
},
3650
};
3751

3852
/// Parse a Parquet file into an array of `T`, reading it as Arrow.
@@ -128,6 +142,203 @@ fn parquet_input_test(compression: Compression) {
128142
}
129143
}
130144

145+
/// The `parquet` input format decodes INT96 timestamps without going through
146+
/// nanoseconds, so a far-future or far-past date survives ingest.
147+
#[test]
148+
fn parquet_input_int96_timestamps() {
149+
let test_data = timestamp_test_data();
150+
let temp_file = NamedTempFile::new().unwrap();
151+
write_timestamp_parquet(temp_file.path(), &test_data, TimestampEncoding::Int96);
152+
153+
let config = serde_json::from_value(json!({
154+
"stream": "test_input",
155+
"transport": {
156+
"name": "file_input",
157+
"config": {
158+
"path": temp_file.path()
159+
}
160+
},
161+
"format": {
162+
"name": "parquet"
163+
}
164+
}))
165+
.unwrap();
166+
167+
let (endpoint, _consumer, _parser, zset) =
168+
mock_input_pipeline::<TimestampTestStruct, TimestampTestStruct>(
169+
config,
170+
Relation::new(
171+
"test".into(),
172+
TimestampTestStruct::schema(),
173+
false,
174+
BTreeMap::new(),
175+
),
176+
)
177+
.unwrap();
178+
endpoint.extend();
179+
wait(
180+
|| {
181+
endpoint.queue(false);
182+
zset.state().flushed.len() == test_data.len()
183+
},
184+
DEFAULT_TIMEOUT_MS,
185+
)
186+
.unwrap();
187+
188+
for (i, upd) in zset.state().flushed.iter().enumerate() {
189+
assert_eq!(upd.unwrap_insert(), &test_data[i]);
190+
}
191+
}
192+
193+
/// Sub-microsecond precision in an INT96 timestamp is dropped, not rounded.
194+
///
195+
/// INT96 counts nanoseconds within the day, so a writer can store a remainder
196+
/// finer than Feldera's `TIMESTAMP` holds. `Int96::to_micros` truncates it
197+
/// toward zero, which is a floor even before the epoch, where the negative
198+
/// part sits in the day number rather than the nanoseconds.
199+
///
200+
/// Each row below reads back as exactly the microsecond it was built from, and
201+
/// every remainder is large enough that rounding would give the next one.
202+
#[test]
203+
fn parquet_input_int96_truncates_sub_microseconds() {
204+
fn at(rfc3339: &str) -> Timestamp {
205+
Timestamp::from_dateTime(DateTime::parse_from_rfc3339(rfc3339).unwrap().to_utc())
206+
}
207+
208+
let rows = [
209+
(at("2024-06-01T12:00:00.123456Z"), 789),
210+
// Before the epoch, so the day number is negative while the
211+
// nanoseconds within the day stay positive.
212+
(at("1600-01-01T00:00:00.123456Z"), 789),
213+
// Past the nanosecond range, and one nanosecond below a whole
214+
// microsecond.
215+
(at("4000-12-31T00:00:00.999999Z"), 999),
216+
];
217+
218+
let temp_file = NamedTempFile::new().unwrap();
219+
write_sub_microsecond_int96_parquet(temp_file.path(), &rows);
220+
221+
let data = Bytes::from(std::fs::read(temp_file.path()).unwrap());
222+
let batches: Vec<RecordBatch> = open_parquet_reader(data, 1024)
223+
.unwrap()
224+
.map(|batch| batch.unwrap())
225+
.collect();
226+
let column = batches[0]
227+
.column(0)
228+
.as_any()
229+
.downcast_ref::<TimestampMicrosecondArray>()
230+
.expect("INT96 must decode at microsecond resolution");
231+
232+
let read_back: Vec<i64> = column.values().to_vec();
233+
let expected: Vec<i64> = rows.iter().map(|(ts, _)| ts.microseconds()).collect();
234+
assert_eq!(read_back, expected);
235+
}
236+
237+
/// A file whose two timestamp columns carry different physical encodings
238+
/// decodes each one on its own terms: the INT96 column reads back retyped to
239+
/// microseconds, the INT64 microsecond column exactly as it was written, down
240+
/// to its UTC marker.
241+
///
242+
/// Coercion keys off a column's own physical type, so neither column's
243+
/// treatment depends on the other. Spark leaves a file in this shape when it
244+
/// rewrites some timestamp columns and not others.
245+
#[test]
246+
fn parquet_input_mixed_timestamp_columns_in_one_file() {
247+
let test_data = timestamp_test_data();
248+
let temp_file = NamedTempFile::new().unwrap();
249+
write_mixed_column_parquet(temp_file.path(), &test_data);
250+
251+
let data = Bytes::from(std::fs::read(temp_file.path()).unwrap());
252+
let batches: Vec<RecordBatch> = open_parquet_reader(data, 1024)
253+
.unwrap()
254+
.map(|batch| batch.unwrap())
255+
.collect();
256+
let batch = &batches[0];
257+
258+
// INT96 carries no time zone, so coercion lands on a naive microsecond
259+
// timestamp; the INT64 column keeps the UTC marker it was written with.
260+
let schema = batch.schema();
261+
assert_eq!(
262+
schema.field_with_name("ts_int96").unwrap().data_type(),
263+
&ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
264+
);
265+
assert_eq!(
266+
schema.field_with_name("ts_micros").unwrap().data_type(),
267+
&ArrowDataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
268+
);
269+
270+
let expected: Vec<Option<i64>> = test_data
271+
.iter()
272+
.map(|row| row.ts.map(|ts| ts.microseconds()))
273+
.collect();
274+
for name in ["ts_int96", "ts_micros"] {
275+
let column = batch
276+
.column_by_name(name)
277+
.unwrap()
278+
.as_any()
279+
.downcast_ref::<TimestampMicrosecondArray>()
280+
.unwrap_or_else(|| panic!("{name} must decode at microsecond resolution"));
281+
let read_back: Vec<Option<i64>> = column.iter().collect();
282+
assert_eq!(read_back, expected, "column {name}");
283+
}
284+
}
285+
286+
/// An INT96 timestamp inside a struct field and inside a list element decodes
287+
/// at microsecond resolution, the same as one at the top level.
288+
///
289+
/// Delta permits timestamps at any depth and Spark writes those as INT96 too.
290+
/// Coercion reaches them by walking struct, list and map paths, including the
291+
/// `.list.element` name Parquet gives a list's elements.
292+
#[test]
293+
fn parquet_input_nested_int96_timestamps() {
294+
let test_data = timestamp_test_data();
295+
let temp_file = NamedTempFile::new().unwrap();
296+
write_nested_int96_parquet(temp_file.path(), &test_data);
297+
298+
let data = Bytes::from(std::fs::read(temp_file.path()).unwrap());
299+
let batches: Vec<RecordBatch> = open_parquet_reader(data, 1024)
300+
.unwrap()
301+
.map(|batch| batch.unwrap())
302+
.collect();
303+
let batch = &batches[0];
304+
305+
// The fixture skips rows without a timestamp.
306+
let expected: Vec<i64> = test_data
307+
.iter()
308+
.filter_map(|row| row.ts)
309+
.map(|ts| ts.microseconds())
310+
.collect();
311+
312+
let in_struct = batch
313+
.column_by_name("nested")
314+
.unwrap()
315+
.as_any()
316+
.downcast_ref::<StructArray>()
317+
.unwrap()
318+
.column_by_name("ts")
319+
.unwrap()
320+
.as_any()
321+
.downcast_ref::<TimestampMicrosecondArray>()
322+
.expect("INT96 inside a struct must decode at microsecond resolution")
323+
.values()
324+
.to_vec();
325+
assert_eq!(in_struct, expected, "struct field");
326+
327+
let in_list = batch
328+
.column_by_name("ts_list")
329+
.unwrap()
330+
.as_any()
331+
.downcast_ref::<ListArray>()
332+
.unwrap()
333+
.values()
334+
.as_any()
335+
.downcast_ref::<TimestampMicrosecondArray>()
336+
.expect("INT96 inside a list must decode at microsecond resolution")
337+
.values()
338+
.to_vec();
339+
assert_eq!(in_list, expected, "list element");
340+
}
341+
131342
#[test]
132343
fn parquet_input_uncompressed() {
133344
parquet_input_test(Compression::UNCOMPRESSED);

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

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
//! thus fully applied (deleted rows are never emitted), and memory stays bounded
1111
//! to one batch.
1212
13+
use crate::format::parquet::int96_as_micros;
1314
use crate::integrated::delta_table::ReadSchema;
1415
use anyhow::{Result as AnyResult, anyhow};
1516
use arrow::array::{
@@ -677,12 +678,17 @@ impl PartitionStream for MaskedParquetPartition {
677678
// decoder builds 32-bit-offset arrays, which cap one batch at 2 GiB
678679
// per column, and `project_to_logical` then pays to cast them.
679680
// Rebuilding the reader metadata reuses the footer read above.
680-
let metadata = ArrowReaderMetadata::try_new(
681-
Arc::clone(probe.metadata()),
682-
ArrowReaderOptions::new()
683-
.with_schema(Arc::clone(ReadSchema::new(probe.schema()).schema())))
681+
let options = ArrowReaderOptions::new()
682+
.with_schema(Arc::clone(ReadSchema::new(probe.schema()).schema()));
683+
let mut metadata = ArrowReaderMetadata::try_new(
684+
Arc::clone(probe.metadata()), options.clone())
684685
.map_err(|e| DataFusionError::External(
685686
format!("failed to read Parquet file '{path}' as view types: {e}").into()))?;
687+
if let Some(coerced) = int96_as_micros(&metadata, &options)
688+
.map_err(|e| DataFusionError::External(
689+
format!("failed to read INT96 timestamps in Parquet file '{path}': {e}").into()))? {
690+
metadata = coerced;
691+
}
686692
// `num_rows()` is `i64` because Parquet's metadata is signed
687693
// throughout; it is non-negative for any file whose footer parsed
688694
// (which it did, just above).

0 commit comments

Comments
 (0)