dwsmith1983 commented on code in PR #5365: URL: https://github.com/apache/datafusion-comet/pull/5365#discussion_r4006991954
########## native/core/src/parquet/datetime_rebase.rs: ########## @@ -0,0 +1,2662 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Per-file datetime calendar-rebase handling for the parquet scan. +//! +//! Spark 2.4 and earlier wrote dates and timestamps in the hybrid Julian + Gregorian calendar; +//! Spark 3.0+ uses the proleptic Gregorian calendar and records the calendar policy of every +//! file it writes in the parquet footer's key-value metadata (`org.apache.spark.version`, +//! `org.apache.spark.legacyDateTime`, `org.apache.spark.legacyINT96`, +//! `org.apache.spark.timeZone`). Spark's reader resolves the rebase policy from EACH FILE's +//! writer metadata (`DataSourceUtils.datetimeRebaseSpec` / `int96RebaseSpec`) -- the session's +//! `spark.sql.parquet.datetimeRebaseModeInRead` conf only applies to files whose metadata does +//! not decide the policy on its own -- so a reader that ignores the metadata silently returns +//! values shifted by up to ten days for dates before 1582-10-15 (e.g. `1500-01-01` reads as +//! `1500-01-10`). +//! +//! This module mirrors that per-file resolution: [`resolve_file_rebase_policies`] computes the +//! date / INT64-timestamp / INT96-timestamp policies from a file's arrow schema metadata (the +//! parquet key-value pairs survive the parquet -> arrow schema conversion), and +//! [`wrap_datetime_rebase`] wraps the per-file rewritten expressions' column references in a +//! [`SparkDatetimeRebaseExpr`] that rebases values exactly where that is possible without the +//! JVM's historical timezone tables (dates always; timestamps for a fixed UTC writer zone) and +//! refuses -- rather than silently corrupting -- ancient values it cannot rebase. Nested +//! columns are rebuilt leaf by leaf (struct / list / map / fixed-size list / dictionary), each +//! leaf under its own policy, with nulls and offsets preserved. Modern values are always the +//! identity under every policy: from 1582-10-15 onward for dates, and from +//! [`LAST_SWITCH_JULIAN_TS_SECONDS`] (1900-01-01T00:00:00Z, Spark's +//! `RebaseDateTime.lastSwitchJulianTs`) onward for timestamps. +//! +//! Spark applies `datetimeRebaseSpec` to INT64 `TIMESTAMP_MICROS` / `TIMESTAMP_MILLIS` columns +//! and `int96RebaseSpec` to INT96 columns. The two physical types are indistinguishable in the +//! arrow schema DataFusion hands the expression adapter (both surface as `Timestamp(us, "UTC")` +//! after INT96 coercion), so Comet's parquet reader factory stamps the file's INT96 leaf +//! ordinals -- taken from the parquet footer's own `SchemaDescriptor` -- into the key-value +//! metadata under [`INT96_LEAVES_METADATA_KEY`] before the arrow schema is derived (see +//! [`stamp_int96_leaves`] and `eager_page_index_reader_factory.rs`), and the adapter attributes +//! every timestamp leaf to its spec from that stamp. Without a stamp, the two specs are merged: +//! agreement decides, disagreement degrades to [`RebasePolicy::CheckAncient`]. +//! +//! The wrapper sits BENEATH the schema adapter's nested narrowing (the struct -> struct convert +//! that keeps only the requested children), which is what keeps those ordinals physical -- but +//! it means the wrapper sees every physical child, requested or not. Spark only ever decodes +//! the requested nested schema, so [`FileRebasePolicies::restrict_to_requested`] marks the +//! physical leaves the narrowing drops as the identity: an unrequested ancient `s.ts` never +//! blocks `select s.d`, exactly as in Spark. +//! +//! Currently only enabled by the Delta scan arms via +//! `SparkParquetOptions::rebase_from_file_metadata`, which also carries the session read modes +//! ([`SessionRebaseModes`], forwarded from the JVM) that decide the policy for files without +//! Spark writer metadata; the plain NativeScan keeps its documented no-rebase behavior (see +//! the compatibility guide and issue #5010). + +use std::collections::HashMap; +use std::fmt::{self, Display}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, Date32Array, FixedSizeListArray, GenericListArray, MapArray, + OffsetSizeTrait, PrimitiveArray, RecordBatch, StructArray, +}; +use arrow::datatypes::{ + ArrowTimestampType, DataType, Date32Type, FieldRef, Schema, SchemaRef, TimeUnit, + TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, + TimestampSecondType, +}; +use arrow::error::ArrowError; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{DataFusionError, Result as DataFusionResult}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ColumnarValue; +use parquet::basic::Type as ParquetPhysicalType; +use parquet::file::metadata::{FileMetaData, KeyValue, ParquetMetaData}; +use parquet::schema::types::SchemaDescriptor; + +use super::name_fold::fold_names; +use super::schema_adapter::parse_field_id; + +/// Footer key naming the Spark release that wrote the file; absent for non-Spark writers. +const SPARK_VERSION_METADATA_KEY: &str = "org.apache.spark.version"; +/// Present (empty value) when the file's dates and INT64 timestamps were written with +/// `spark.sql.parquet.datetimeRebaseModeInWrite=LEGACY`. +const SPARK_LEGACY_DATETIME_KEY: &str = "org.apache.spark.legacyDateTime"; +/// Present (empty value) when the file's INT96 timestamps were written with +/// `spark.sql.parquet.int96RebaseModeInWrite=LEGACY`. +const SPARK_LEGACY_INT96_KEY: &str = "org.apache.spark.legacyINT96"; +/// The writer session's time zone, stamped alongside either legacy flag. +const SPARK_TIMEZONE_KEY: &str = "org.apache.spark.timeZone"; + +/// Key-value metadata entry Comet's parquet reader factory adds to a file's footer metadata +/// (in memory only, never written back) so the expression adapter can tell INT96 timestamp +/// columns from INT64 ones after both have been coerced to the same arrow type. Value: +/// `"<leaf count>:<comma-separated INT96 leaf ordinals>"`, where leaves are the file's +/// primitive columns in `SchemaDescriptor::columns()` order -- the same depth-first order +/// parquet-rs assigns arrow leaves, so an arrow-side depth-first walk lines up with it. The +/// leaf count lets the reader detect a stamp that does not describe the schema it is paired +/// with (see [`Int96Attribution::from_schema`]). +pub(crate) const INT96_LEAVES_METADATA_KEY: &str = "comet.int96_leaf_columns"; + +/// Day of the Gregorian cutover (1582-10-15) as days since the epoch; rebasing is the identity +/// from this day onward. Same value as Spark's `RebaseDateTime.lastSwitchJulianDay`. +const LAST_SWITCH_JULIAN_DAY: i32 = -141427; + +/// Spark's `RebaseDateTime.lastSwitchJulianTs` (and `lastSwitchGregorianTs`) in seconds since +/// the epoch: 1900-01-01T00:00:00Z. Spark derives it as the latest switch instant across every +/// zone in its `julian-gregorian-rebase-micros.json` table (`getLastSwitchTs`, which also +/// asserts the calendars' difference is zero for every zone from then on): most zones ran on +/// local mean time before 1900, so the last instant at which rebasing changes a value in ANY +/// zone is 1900-01-01T00:00:00Z, not the 1582 cutover. `createTimestampRebaseFuncInRead` +/// under `EXCEPTION` throws exactly for `micros < lastSwitchJulianTs` (after converting +/// `TIMESTAMP_MILLIS` to micros), and `rebaseJulianToGregorianMicros` is the identity from it +/// onward in every zone. The value is in seconds so it scales exactly to any timestamp unit. +pub(crate) const LAST_SWITCH_JULIAN_TS_SECONDS: i64 = -2_208_988_800; + +/// The per-century differences between the Julian and proleptic Gregorian calendars, and the +/// Julian-calendar switch days at which each difference starts to apply. Copied verbatim from +/// Spark's `RebaseDateTime.julianGregDiffs` / `julianGregDiffSwitchDay` (which Spark generated +/// from `localRebaseJulianToGregorianDays`); `rebase_julian_to_gregorian_days` must stay +/// value-for-value equal to Spark's `rebaseJulianToGregorianDays`. +const JULIAN_GREG_DIFFS: [i32; 14] = [2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, 0]; +const JULIAN_GREG_DIFF_SWITCH_DAY: [i32; 14] = [ + -719164, -682945, -646420, -609895, -536845, -500320, -463795, -390745, -354220, -317695, + -244645, -208120, -171595, -141427, +]; + +/// Proleptic-Gregorian days since 1970-01-01 for a nominal civil date, via Howard Hinnant's +/// `days_from_civil`. `d` may exceed the month's length; the excess rolls into the following +/// month exactly like `LocalDate.of(y, m, 1).plusDays(d - 1)` in Spark's +/// `localRebaseJulianToGregorianDays` (how the non-existent proleptic date `1000-02-29`, +/// valid in the Julian calendar, lands on `1000-03-01`). +fn days_from_civil(y: i64, m: i64, d: i64) -> i64 { + let y = if m <= 2 { y - 1 } else { y }; + let era = y.div_euclid(400); + let yoe = y - era * 400; // [0, 399] + let mp = (m + 9) % 12; // [0, 11], March = 0 + let doy = (153 * mp + 2) / 5 + d - 1; + let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; + era * 146097 + doe - 719468 +} + +/// Julian-calendar civil date `(year, month, day)` for a day count since 1970-01-01 that labels +/// days in the Julian calendar (astronomical year numbering: 1 BCE is year 0). Standard +/// Julian-day-number conversion (E.G. Richards' algorithm), exact for any day. +fn julian_day_to_civil(days: i64) -> (i64, i64, i64) { + // Integer (noon) Julian Day Number of this civil day: 1970-01-01 is JDN 2440588. + let jdn = days + 2_440_588; + let f = jdn + 1401; + let e = 4 * f + 3; + let g = e.rem_euclid(1461) / 4; + let h = 5 * g + 2; + let day = h.rem_euclid(153) / 5 + 1; + let month = (h / 153 + 2).rem_euclid(12) + 1; + let year = e.div_euclid(1461) - 4716 + (14 - month) / 12; + (year, month, day) +} + +/// Exact port of Spark's `RebaseDateTime.rebaseJulianToGregorianDays`: reinterprets a day count +/// written in the hybrid Julian + Gregorian calendar as the proleptic Gregorian day count of the +/// same nominal civil date. Identity for days from 1582-10-15 onward. Days before the tables' +/// range (before Julian `0001-01-01`) take the calendar-arithmetic path, mirroring Spark's +/// `localRebaseJulianToGregorianDays` fallback. +pub(crate) fn rebase_julian_to_gregorian_days(days: i32) -> i32 { + if days < JULIAN_GREG_DIFF_SWITCH_DAY[0] { + let (y, m, d) = julian_day_to_civil(days as i64); + (days_from_civil(y, m, 1) + (d - 1)) as i32 + } else { + // Spark's rebaseDays: linear search from the most recent switch day. + let mut i = JULIAN_GREG_DIFF_SWITCH_DAY.len(); + loop { + i -= 1; + if i == 0 || days >= JULIAN_GREG_DIFF_SWITCH_DAY[i] { + break; + } + } + days + JULIAN_GREG_DIFFS[i] + } +} + +/// Timezone strings from `org.apache.spark.timeZone` that denote a fixed zero-offset zone in +/// both `java.util.TimeZone` and `java.time`. Only for these is timestamp rebasing the pure +/// nominal-date shift [`SparkDatetimeRebaseExpr::rebase_timestamp_utc`] computes; any other (or +/// absent) zone needs the JVM's historical timezone tables and stays on the +/// refuse-ancient-values path. +const UTC_EQUIVALENT_TIMEZONES: [&str; 6] = ["UTC", "Etc/UTC", "GMT", "Etc/GMT", "Z", "+00:00"]; + +/// How the writer's session time zone (if recorded) affects timestamp rebasing. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum WriterTimeZone { + /// A fixed zero-offset zone: rebasing reduces to the exact nominal-date shift. + Utc, + /// Any other zone, or none recorded (pre-3.0 files): ancient values cannot be rebased + /// without the JVM's historical timezone data. + OtherOrUnknown, +} + +/// One session-level datetime rebase read mode (a `LegacyBehaviorPolicy` value of +/// `spark.sql.parquet.datetimeRebaseModeInRead` / `int96RebaseModeInRead`), consulted by +/// [`resolve_file_rebase_policies`] ONLY for files whose footer metadata does not decide the +/// policy on its own -- exactly the `getOrElse` fallback in Spark's +/// `DataSourceUtils.getRebaseSpec`. Files that carry `org.apache.spark.version` ignore these +/// modes entirely, on every Spark version. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub(crate) enum RebaseReadMode { + /// Refuse ancient values (Spark raises `SparkUpgradeException`); maps to + /// [`RebasePolicy::CheckAncient`]. The default mirrors the conservative posture used + /// before the conf was plumbed through (and Spark 3.x's own conf default). + #[default] + Exception, + /// Read values as proleptic Gregorian without rebasing. + Corrected, + /// Rebase from the hybrid Julian + Gregorian calendar. + Legacy, +} + +impl RebaseReadMode { + /// Parses a `LegacyBehaviorPolicy` conf value. `SQLConf` validates and upper-cases the + /// session conf, but a per-relation `datetimeRebaseMode` option arrives verbatim, so the + /// match is case-insensitive. Anything unrecognized -- including the empty string a proto + /// producer that predates the field sends -- falls back to [`RebaseReadMode::Exception`], + /// which refuses ancient values rather than silently corrupting them. + pub(crate) fn from_conf_value(value: &str) -> Self { + match value.to_ascii_uppercase().as_str() { + "CORRECTED" => RebaseReadMode::Corrected, + "LEGACY" => RebaseReadMode::Legacy, + _ => RebaseReadMode::Exception, + } + } +} + +/// The session's effective datetime rebase read modes, one per spec class (INT64 +/// dates/timestamps vs INT96 timestamps), forwarded from the JVM at planning time. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub(crate) struct SessionRebaseModes { + /// `spark.sql.parquet.datetimeRebaseModeInRead` (or the relation's `datetimeRebaseMode`). + pub datetime: RebaseReadMode, + /// `spark.sql.parquet.int96RebaseModeInRead` (or the relation's `int96RebaseMode`). + pub int96: RebaseReadMode, +} + +/// Calendar policy of one file's date or timestamp columns, resolved from writer metadata the +/// same way Spark's `DataSourceUtils.getRebaseSpec` resolves it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum RebasePolicy { + /// Written in the proleptic Gregorian calendar; values pass through untouched. + Corrected, + /// Written in the hybrid Julian + Gregorian calendar; values must be rebased. + Legacy(WriterTimeZone), + /// Policy could not be pinned down (contradictory flags, or a non-Spark writer under the + /// `EXCEPTION` read mode): modern values -- identical under either calendar -- pass, + /// ancient values raise. Mirrors Spark's `EXCEPTION` behavior (`SparkUpgradeException`). + CheckAncient, +} + +/// Which of a file's leaf columns are physically INT96, from the stamp the parquet reader +/// factory adds under [`INT96_LEAVES_METADATA_KEY`]. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) enum Int96Attribution { + /// No stamp, or a stamp whose leaf count does not match the schema it arrived with: the + /// INT64 and INT96 timestamp specs cannot be told apart per column and are merged. + Unknown, + /// Sorted leaf ordinals (depth-first over the file schema's primitive columns) that are + /// INT96; every other timestamp leaf is INT64. + Known(Vec<usize>), +} + +impl Int96Attribution { + /// Parses the stamp out of `schema`'s metadata and validates its leaf count against the + /// schema's own depth-first leaf count, so a stamp that does not describe this schema (a + /// crafted footer key, or a cached-metadata mismatch) degrades to [`Self::Unknown`]. + fn from_schema(schema: &Schema) -> Self { + let Some(stamp) = schema.metadata().get(INT96_LEAVES_METADATA_KEY) else { + return Int96Attribution::Unknown; + }; + let Some((count, ordinals)) = stamp.split_once(':') else { + return Int96Attribution::Unknown; + }; + let schema_leaves: usize = schema + .fields() + .iter() + .map(|f| leaf_count(f.data_type())) + .sum(); + if count.parse::<usize>().ok() != Some(schema_leaves) { + return Int96Attribution::Unknown; + } + let parsed: Option<Vec<usize>> = if ordinals.is_empty() { + Some(Vec::new()) + } else { + ordinals + .split(',') + .map(|o| o.parse::<usize>().ok().filter(|o| *o < schema_leaves)) + .collect() + }; + match parsed { + Some(mut leaves) => { + leaves.sort_unstable(); + Int96Attribution::Known(leaves) + } + None => Int96Attribution::Unknown, + } + } + + /// `Some(true)` / `Some(false)` when the leaf is known to be INT96 / INT64, `None` when + /// the attribution is unknown. + fn is_int96(&self, leaf: usize) -> Option<bool> { + match self { + Int96Attribution::Unknown => None, + Int96Attribution::Known(leaves) => Some(leaves.binary_search(&leaf).is_ok()), + } + } +} + +/// The [`INT96_LEAVES_METADATA_KEY`] value describing `schema`: its leaf count and the +/// ordinals of its INT96 primitive columns. +pub(crate) fn int96_leaf_stamp(schema: &SchemaDescriptor) -> String { + let ordinals: Vec<String> = schema + .columns() + .iter() + .enumerate() + .filter(|(_, column)| column.physical_type() == ParquetPhysicalType::INT96) + .map(|(ordinal, _)| ordinal.to_string()) + .collect(); + format!("{}:{}", schema.num_columns(), ordinals.join(",")) +} + +/// Returns a copy of `metadata` whose key-value metadata carries the [`int96_leaf_stamp`] of +/// its own schema, or `None` when it already does (the common case after the first open of a +/// file, since the caller caches the stamped copy). Any pre-existing entry under the key -- +/// a file cannot legitimately carry one -- is replaced, never trusted. Only the file-level +/// key-value list changes; row groups and page indexes are carried over as-is. The parquet +/// API cannot carry a file decryptor, nor `FileMetaData`'s crate-private encryption fields +/// (encryption algorithm, footer signing key metadata), across this rebuild, so callers must +/// not stamp opens that supply decryption properties -- and the only consumer, the Delta +/// scan, declines every encrypted-parquet configuration before planning, so a parquet +/// modular encryption file never reaches this path with or without those properties. +pub(crate) fn stamp_int96_leaves(metadata: &ParquetMetaData) -> Option<ParquetMetaData> { + let file_metadata = metadata.file_metadata(); + let stamp = int96_leaf_stamp(file_metadata.schema_descr()); + let existing = file_metadata + .key_value_metadata() + .and_then(|kvs| kvs.iter().find(|kv| kv.key == INT96_LEAVES_METADATA_KEY)) + .and_then(|kv| kv.value.as_deref()); + if existing == Some(stamp.as_str()) { + return None; + } + let mut key_values: Vec<KeyValue> = file_metadata + .key_value_metadata() + .map(|kvs| { + kvs.iter() + .filter(|kv| kv.key != INT96_LEAVES_METADATA_KEY) + .cloned() + .collect() + }) + .unwrap_or_default(); + key_values.push(KeyValue::new(INT96_LEAVES_METADATA_KEY.to_string(), stamp)); + let stamped_file_metadata = FileMetaData::new( + file_metadata.version(), + file_metadata.num_rows(), + file_metadata.created_by().map(str::to_string), + Some(key_values), + file_metadata.schema_descr_ptr(), + file_metadata.column_orders().cloned(), + ); + Some( + ParquetMetaData::new(stamped_file_metadata, metadata.row_groups().to_vec()) + .into_builder() + .set_column_index(metadata.column_index().cloned()) + .set_offset_index(metadata.offset_index().cloned()) + .build(), + ) +} + +/// Per-file rebase policies for the three affected column classes, plus the INT96 +/// attribution that selects between the two timestamp specs per leaf. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct FileRebasePolicies { + /// `DATE` columns, governed by `org.apache.spark.legacyDateTime` alone. + pub date: RebasePolicy, + /// INT64 `TIMESTAMP_MICROS` / `TIMESTAMP_MILLIS` columns: the datetime spec (same + /// resolution as `date`), as Spark's `ParquetVectorUpdaterFactory` selects for INT64. + pub int64_timestamp: RebasePolicy, + /// INT96 columns: the INT96 spec (`org.apache.spark.legacyINT96`, min version 3.1.0). + pub int96_timestamp: RebasePolicy, + /// Which timestamp leaves are INT96. See [`Int96Attribution`]. + pub int96_leaves: Int96Attribution, + /// Sorted depth-first leaf ordinals -- over the physical file schema, the same ordinals + /// `int96_leaves` uses -- that the query does not read: nested children the schema + /// adapter's struct narrowing drops before any value leaves the scan. Spark never decodes + /// them either, so their policy is the identity whatever the file's calendar. Empty until + /// [`Self::restrict_to_requested`] runs (every leaf requested). + pub unrequested_leaves: Vec<usize>, +} + +impl FileRebasePolicies { + /// True when some policy is not the plain proleptic-Gregorian pass-through, i.e. when the + /// per-column wrap in [`wrap_datetime_rebase`] can install anything at all. + pub(crate) fn any_rebase_needed(&self) -> bool { + self.date != RebasePolicy::Corrected + || self.int64_timestamp != RebasePolicy::Corrected + || self.int96_timestamp != RebasePolicy::Corrected + } + + fn is_requested(&self, leaf: usize) -> bool { + self.unrequested_leaves.binary_search(&leaf).is_err() + } + + /// The policy of the `Date32` leaf at depth-first ordinal `leaf`: the file's date policy, + /// or the identity when the query does not read that leaf. + fn date_policy(&self, leaf: usize) -> RebasePolicy { + if self.is_requested(leaf) { + self.date + } else { + RebasePolicy::Corrected + } + } + + /// The policy of the timezone-carrying timestamp leaf at depth-first ordinal `leaf`: the + /// identity when the query does not read it; otherwise its physical type's spec when the + /// attribution is known, or else the two specs merged -- agreement decides, disagreement + /// degrades to [`RebasePolicy::CheckAncient`], which still passes every modern value and + /// refuses only ancient ones. + fn timestamp_policy(&self, leaf: usize) -> RebasePolicy { + if !self.is_requested(leaf) { + return RebasePolicy::Corrected; + } + match self.int96_leaves.is_int96(leaf) { + Some(true) => self.int96_timestamp, + Some(false) => self.int64_timestamp, + None if self.int64_timestamp == self.int96_timestamp => self.int64_timestamp, + None => RebasePolicy::CheckAncient, + } + } + + /// These policies with every physical leaf the query does not read marked the identity. + /// `requested` pairs each top-level field of `physical_schema` (by position) with the type + /// of the logical field the schema adapter narrows it to -- `None` for a column without a + /// logical counterpart, whose leaves are left as they are (no expression reads it anyway). + /// Nested children pair the way the adapter's struct convert selects them (see + /// [`push_unrequested_leaves`]); the INT96 attribution is untouched, since the ordinals + /// stay physical. `requested` is parallel to the schema's fields; should a caller pass a + /// shorter slice, the trailing columns simply keep every leaf (the safe direction). + pub(crate) fn restrict_to_requested( + mut self, + physical_schema: &Schema, + requested: &[Option<&DataType>], + case_sensitive: bool, + use_field_id: bool, + ) -> Self { + debug_assert_eq!(requested.len(), physical_schema.fields().len()); + let matching = FieldMatching { + case_sensitive, + use_field_id, + }; + let mut next_leaf = 0; + let mut unrequested = Vec::new(); + for (field, requested) in physical_schema.fields().iter().zip(requested) { + match requested { + Some(logical) => push_unrequested_leaves( + field.data_type(), + logical, + &mut next_leaf, + matching, + &mut unrequested, + ), + None => next_leaf += leaf_count(field.data_type()), + } + } + // Emitted in depth-first order, so already sorted for `is_requested`'s binary search. + self.unrequested_leaves = unrequested; + self + } +} + +/// The field-matching rules of the schema adapter's nested narrowing +/// (`parquet_convert_struct_to_struct`): names fold per `case_sensitive`, and Parquet field ids +/// select fields when `use_field_id` is set. +#[derive(Debug, Clone, Copy)] +struct FieldMatching { + case_sensitive: bool, + use_field_id: bool, +} + +/// Appends to `out` the depth-first leaf ordinals of `physical` (counting from `next_leaf`, +/// which advances past every leaf of `physical`) that reading it as `requested` drops. +/// +/// Recurses through exactly the pairings `parquet_convert_array` narrows, and no others: a +/// struct child is dropped only when NO requested child selects it by either rule the struct +/// convert uses -- folded name, or Parquet field id when ids are in play -- and an ambiguous +/// child (several requested children select it) is kept; `List` pairs with `List` by element +/// type, and `Map` with a `Map` of the same key ordering by its entries, positionally. Any +/// other pairing -- a leaf, a `LargeList` / `FixedSizeList` / dictionary, a map whose ordering +/// differs, or a shape mismatch -- is handed to arrow's cast or passed through whole by the +/// convert, so it keeps every leaf. Keeping a superset of what the narrowing reads is always +/// safe (a spurious check at worst); dropping a leaf the narrowing reads would skip its +/// rebase, so every doubt resolves to "requested". +fn push_unrequested_leaves( + physical: &DataType, + requested: &DataType, + next_leaf: &mut usize, + matching: FieldMatching, + out: &mut Vec<usize>, +) { + match (physical, requested) { + (DataType::Struct(physical_fields), DataType::Struct(requested_fields)) => { + let names: Vec<&str> = physical_fields + .iter() + .chain(requested_fields.iter()) + .map(|f| f.name().as_str()) + .collect(); + let folded = fold_names(&names, matching.case_sensitive); + let (physical_folded, requested_folded) = folded.split_at(physical_fields.len()); + for (i, child) in physical_fields.iter().enumerate() { + let child_id = if matching.use_field_id { + parse_field_id(child) + } else { + None + }; + let mut selectors = requested_fields.iter().enumerate().filter(|(j, r)| { + requested_folded[*j] == physical_folded[i] + || (child_id.is_some() && parse_field_id(r) == child_id) + }); + match (selectors.next(), selectors.next()) { + (None, _) => { + let n = leaf_count(child.data_type()); + out.extend(*next_leaf..*next_leaf + n); + *next_leaf += n; + } + (Some((_, requested_child)), None) => push_unrequested_leaves( + child.data_type(), + requested_child.data_type(), + next_leaf, + matching, + out, + ), + (Some(_), Some(_)) => *next_leaf += leaf_count(child.data_type()), + } + } + } + (DataType::List(physical_item), DataType::List(requested_item)) => push_unrequested_leaves( + physical_item.data_type(), + requested_item.data_type(), + next_leaf, + matching, + out, + ), + ( + DataType::Map(physical_entries, physical_sorted), + DataType::Map(requested_entries, requested_sorted), + ) if physical_sorted == requested_sorted => { + match (physical_entries.data_type(), requested_entries.data_type()) { + (DataType::Struct(physical_kv), DataType::Struct(requested_kv)) + if physical_kv.len() == requested_kv.len() => + { + for (p, r) in physical_kv.iter().zip(requested_kv.iter()) { + push_unrequested_leaves( + p.data_type(), + r.data_type(), + next_leaf, + matching, + out, + ); + } + } + _ => *next_leaf += leaf_count(physical), + } + } + _ => *next_leaf += leaf_count(physical), + } +} + +/// The writer time zone recorded in `metadata`, classified for timestamp rebasing. Mirrors the +/// `Option(lookupFileMeta(SPARK_TIMEZONE_METADATA_KEY))` lookup Spark's `getRebaseSpec` performs +/// for every LEGACY resolution, conf-fallback included; Spark substitutes the JVM default zone +/// when the key is absent (`RebaseSpec.timeZone`), which is unavailable natively, so an absent or +/// non-UTC zone classifies as [`WriterTimeZone::OtherOrUnknown`] (dates still rebase fully -- +/// the day rebase is zone-free -- while ancient timestamps refuse rather than guess). +fn writer_time_zone(metadata: &HashMap<String, String>) -> WriterTimeZone { + match metadata.get(SPARK_TIMEZONE_KEY) { + Some(tz) if UTC_EQUIVALENT_TIMEZONES.contains(&tz.as_str()) => WriterTimeZone::Utc, + _ => WriterTimeZone::OtherOrUnknown, + } +} + +/// One spec resolution, mirroring Spark's `DataSourceUtils.getRebaseSpec` exactly: a Spark +/// version below `min_version` (lexicographic comparison, same as the Scala `String.<`) or a +/// present legacy flag means LEGACY; a Spark version at/after `min_version` without the flag +/// means CORRECTED; no Spark version at all falls back to `conf_mode`, the session read conf +/// forwarded from the JVM (`getRebaseSpec`'s `modeByConfig` fallback, its ONLY use of the +/// conf): CORRECTED passes values through, LEGACY rebases (with the writer zone from the +/// file's `org.apache.spark.timeZone` key, same lookup as the metadata-driven LEGACY path), +/// and EXCEPTION refuses ancient values as [`RebasePolicy::CheckAncient`]. +fn resolve_spec( + metadata: &HashMap<String, String>, + min_version: &str, + legacy_key: &str, + conf_mode: RebaseReadMode, +) -> RebasePolicy { + match metadata.get(SPARK_VERSION_METADATA_KEY) { + None => match conf_mode { + RebaseReadMode::Corrected => RebasePolicy::Corrected, + RebaseReadMode::Legacy => RebasePolicy::Legacy(writer_time_zone(metadata)), + RebaseReadMode::Exception => RebasePolicy::CheckAncient, + }, + Some(version) => { + if version.as_str() < min_version || metadata.contains_key(legacy_key) { + RebasePolicy::Legacy(writer_time_zone(metadata)) + } else { + RebasePolicy::Corrected + } + } + } +} + +/// Resolves the per-file rebase policies from a file's arrow schema: the parquet footer's +/// key-value pairs in its metadata decide the specs (the datetime spec uses min version +/// `3.0.0` and the INT96 spec `3.1.0`, matching `DataSourceUtils.datetimeRebaseSpec` / +/// `int96RebaseSpec`; `session_modes` supplies the per-spec conf fallback for files without +/// Spark writer metadata), and the reader factory's INT96 stamp -- validated against the +/// schema's leaf structure -- attributes each timestamp leaf to its spec. +pub(crate) fn resolve_file_rebase_policies( + physical_file_schema: &Schema, + session_modes: SessionRebaseModes, +) -> FileRebasePolicies { + let metadata = physical_file_schema.metadata(); + let datetime_spec = resolve_spec( + metadata, + "3.0.0", + SPARK_LEGACY_DATETIME_KEY, + session_modes.datetime, + ); + let int96_spec = resolve_spec( + metadata, + "3.1.0", + SPARK_LEGACY_INT96_KEY, + session_modes.int96, + ); + FileRebasePolicies { + date: datetime_spec, + int64_timestamp: datetime_spec, + int96_timestamp: int96_spec, + int96_leaves: Int96Attribution::from_schema(physical_file_schema), + unrequested_leaves: Vec::new(), + } +} + +/// Number of primitive leaves `dt` contains in a depth-first walk -- the same count and order +/// parquet-rs uses when it maps the file's `SchemaDescriptor` columns onto the arrow schema, so +/// arrow-side leaf ordinals line up with [`int96_leaf_stamp`]'s. +fn leaf_count(dt: &DataType) -> usize { + match dt { + DataType::Struct(fields) => fields.iter().map(|f| leaf_count(f.data_type())).sum(), + DataType::List(f) + | DataType::LargeList(f) + | DataType::FixedSizeList(f, _) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::Map(f, _) => leaf_count(f.data_type()), + DataType::Dictionary(_, value) => leaf_count(value), + DataType::RunEndEncoded(_, value) => leaf_count(value.data_type()), + DataType::Union(fields, _) => fields.iter().map(|(_, f)| leaf_count(f.data_type())).sum(), + _ => 1, + } +} + +/// Appends the policy of every leaf of `dt`, in depth-first order, to `out`, consuming leaf +/// ordinals from `next_leaf` (exactly [`leaf_count`] of them). Only `Date32` and +/// timezone-carrying timestamps have a policy to apply, and only when the query reads the +/// leaf; timezone-free timestamps are `TIMESTAMP_NTZ`, which Spark never rebases, and every +/// other leaf is the identity ([`RebasePolicy::Corrected`]). +fn leaf_policies( + dt: &DataType, + next_leaf: &mut usize, + policies: &FileRebasePolicies, + out: &mut Vec<RebasePolicy>, +) { + match dt { + DataType::Date32 => { + out.push(policies.date_policy(*next_leaf)); + *next_leaf += 1; + } + DataType::Timestamp(_, Some(_)) => { + out.push(policies.timestamp_policy(*next_leaf)); + *next_leaf += 1; + } + DataType::Struct(fields) => { + for f in fields { + leaf_policies(f.data_type(), next_leaf, policies, out); + } + } + // Mirrors `leaf_count` variant for variant, so a rebase-affected leaf inside a nested + // type `rebase_array` cannot rebuild (views, run-end, union -- never produced from a + // parquet schema) still gets its real policy and makes `rebase_array` refuse loudly + // instead of being stamped the identity. + DataType::List(f) + | DataType::LargeList(f) + | DataType::FixedSizeList(f, _) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::Map(f, _) => leaf_policies(f.data_type(), next_leaf, policies, out), + DataType::Dictionary(_, value) => leaf_policies(value, next_leaf, policies, out), + DataType::RunEndEncoded(_, value) => { + leaf_policies(value.data_type(), next_leaf, policies, out) + } + DataType::Union(fields, _) => { + for (_, f) in fields.iter() { + leaf_policies(f.data_type(), next_leaf, policies, out); + } + } + _ => { + *next_leaf += 1; + out.push(RebasePolicy::Corrected); + } + } +} + +/// Wraps every column reference in `expr` whose physical file type contains a rebase-affected +/// leaf under a policy that needs handling with a [`SparkDatetimeRebaseExpr`] carrying that +/// column's per-leaf policies, so both the per-file projection and the pushed-down predicate +/// evaluate rebased values. Columns whose leaves are all the identity -- unaffected types, +/// affected types under [`RebasePolicy::Corrected`], or leaves the query does not read (see +/// [`FileRebasePolicies::restrict_to_requested`]) -- pass through unwrapped. (The pruning +/// predicates derived from the wrapped predicate treat the wrapper as an opaque expression and +/// skip pruning on those columns -- conservative, since file-level statistics are in the +/// file's own calendar.) +pub(crate) fn wrap_datetime_rebase( + expr: Arc<dyn PhysicalExpr>, + physical_schema: &SchemaRef, + policies: &FileRebasePolicies, +) -> DataFusionResult<Arc<dyn PhysicalExpr>> { + expr.transform(|e| { + let Some(col) = e.downcast_ref::<Column>() else { + return Ok(Transformed::no(e)); + }; + // Missing columns were already replaced with literals; any surviving reference is + // physical-schema-indexed. Out-of-range means a non-file column (defensive): skip. + let Some(field) = physical_schema.fields().get(col.index()) else { + return Ok(Transformed::no(e)); + }; + // This column's first leaf ordinal: the leaves of every preceding top-level field. + let mut next_leaf: usize = physical_schema.fields()[..col.index()] + .iter() + .map(|f| leaf_count(f.data_type())) + .sum(); + let mut column_leaf_policies = Vec::with_capacity(leaf_count(field.data_type())); + leaf_policies( + field.data_type(), + &mut next_leaf, + policies, + &mut column_leaf_policies, + ); + if column_leaf_policies + .iter() + .all(|p| *p == RebasePolicy::Corrected) + { + return Ok(Transformed::no(e)); + } + Ok(Transformed::yes(Arc::new(SparkDatetimeRebaseExpr { + child: e, + field: Arc::clone(field), + leaf_policies: column_leaf_policies, + }) as Arc<dyn PhysicalExpr>)) + }) + .map(|t| t.data) +} + +/// Applies a file's calendar-rebase policies to one column: rebases exactly where possible, +/// raises on ancient values it cannot rebase, and passes modern values (the identity under +/// every policy) through untouched. Nested columns are rebuilt leaf by leaf with nulls and +/// offsets preserved. See the module doc for the policy table. +#[derive(Debug, Eq)] +struct SparkDatetimeRebaseExpr { + child: Arc<dyn PhysicalExpr>, + /// The physical file field this expression reads (type preserved by the rebase). + field: FieldRef, + /// One policy per primitive leaf of `field`'s type, in depth-first order (a single entry + /// for a flat column). At least one is not [`RebasePolicy::Corrected`]. + leaf_policies: Vec<RebasePolicy>, +} + +impl SparkDatetimeRebaseExpr { + /// The refusal error, as an [`ArrowError`] so `try_unary` closures can raise it directly; + /// it converts into a `DataFusionError` at the `?` in `evaluate`. + fn rebase_error(&self, detail: &str) -> ArrowError { + ArrowError::ComputeError(format!( + "Native scan cannot rebase ancient values in column '{}': the file was written \ + with the legacy (hybrid Julian/Gregorian) calendar, or does not declare which \ + calendar it used, and {detail}. Reading it natively would return silently \ + shifted values; disable the native Delta scan \ + (spark.comet.scan.delta.enabled=false) to let Spark read this table", + self.field.name(), + )) + } + + fn internal_error(&self, detail: impl Display) -> DataFusionError { + DataFusionError::Internal(format!( + "SparkDatetimeRebaseExpr on column '{}': {detail}", + self.field.name() + )) + } + + /// Rebases a timestamp column written at a fixed zero-offset zone: shift the nominal day + /// with the exact date table, keep the time of day. Matches Spark's + /// `rebaseJulianToGregorianMicros` for UTC, where the hybrid calendar's day boundaries sit + /// exactly on multiples of a day and no timezone transition can apply (UTC's last switch + /// instant in Spark's rebase table is the 1582-10-15 cutover itself). + fn rebase_timestamp_utc(&self, v: i64, units_per_day: i64) -> Result<i64, ArrowError> { + // Compare in days, not units: the cutover day times a nanosecond day does not fit i64. + let day = v.div_euclid(units_per_day); + if day >= LAST_SWITCH_JULIAN_DAY as i64 { + return Ok(v); + } + let time_of_day = v - day * units_per_day; + let day = i32::try_from(day).map_err(|_| { + self.rebase_error("the value is outside the rebaseable timestamp range") + })?; + let rebased = rebase_julian_to_gregorian_days(day) as i64; + rebased + .checked_mul(units_per_day) + .and_then(|d| d.checked_add(time_of_day)) + .ok_or_else(|| self.rebase_error("the rebased value overflows the timestamp range")) + } + + /// The refuse-ancient-values policy for timestamps: values from + /// [`LAST_SWITCH_JULIAN_TS_SECONDS`] onward are identical under both calendars in every + /// zone (Spark's `createTimestampRebaseFuncInRead` under `EXCEPTION` accepts exactly + /// these, and `rebaseJulianToGregorianMicros` is the identity on them for any zone); + /// older values raise. + fn check_ancient_timestamp( + &self, + v: i64, + units_per_second: i64, + detail: &str, + ) -> Result<i64, ArrowError> { + if v >= LAST_SWITCH_JULIAN_TS_SECONDS * units_per_second { + Ok(v) + } else { + Err(self.rebase_error(detail)) + } + } + + fn rebase_timestamp_array<T: ArrowTimestampType>( + &self, + array: &PrimitiveArray<T>, + policy: RebasePolicy, + units_per_second: i64, + ) -> DataFusionResult<ArrayRef> { + let tz = array.timezone().map(Arc::<str>::from); + let rebased: PrimitiveArray<T> = match policy { + RebasePolicy::Corrected => return Ok(Arc::new(array.clone())), + RebasePolicy::Legacy(WriterTimeZone::Utc) => arrow::compute::try_unary(array, |v| { + self.rebase_timestamp_utc(v, units_per_second * 86_400) + })?, + RebasePolicy::Legacy(WriterTimeZone::OtherOrUnknown) => { Review Comment: > A validity-aware `all(v >= cutoff)` over `values()` followed by `Ok(Arc::clone(array))` would make them allocation-free Done for timestamps and dates. A batch with nothing before the cutover returns the input Arc under every policy, including `Legacy(Utc)`, and the check is validity-aware so null slots holding ancient values do not force a copy. Tests assert pointer equality for the pass-through, that the checking policies still reject an ancient non-null value, and that the legacy UTC rebase still produces a new buffer when one is needed. The `try_unary` helper for the checking arms is gone since nothing calls it now. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
