Skip to content

fix: prevent silent overflow when reading Parquet TIMESTAMP_MILLIS values #5177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
peterxcli wants to merge 8 commits into
base: main
Choose a base branch
from
Open

fix: prevent silent overflow when reading Parquet TIMESTAMP_MILLIS values #5177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a68a595
refactor: use Arrow casts for temporal conversions
peterxcli Jul 31, 2026
263a61e
fix: preserve Spark temporal cast semantics
peterxcli Aug 2, 2026
bd1d243
test: match Spark micros-to-millis cases
peterxcli Aug 2, 2026
fbb2a7c
test: link Spark micros-to-millis cases
peterxcli Aug 2, 2026
47211ff
test: use Spark tag in source link
peterxcli Aug 2, 2026
6afef50
Use imported arity kernel for Spark timestamp downscaling
peterxcli Aug 2, 2026
e71f480
Merge branch 'main' into refactor/5090-use-arrow-temporal-casts
peterxcli Aug 2, 2026
88685aa
fix: enforce Spark temporal conversion semantics
peterxcli Aug 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
262 changes: 51 additions & 211 deletions native/core/src/parquet/cast_column.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,15 @@
// specific language governing permissions and limitations
// under the License.
use arrow::{
array::{
make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray,
TimestampMicrosecondArray, TimestampMillisecondArray,
},
array::{make_array, Array, ArrayRef, LargeListArray, ListArray, MapArray, StructArray},
compute::CastOptions,
datatypes::{DataType, FieldRef, Schema, TimeUnit},
record_batch::RecordBatch,
};

use crate::parquet::parquet_support::{spark_parquet_convert, SparkParquetOptions};
use datafusion::common::format::DEFAULT_CAST_OPTIONS;
use datafusion::common::Result as DataFusionResult;
use datafusion::common::ScalarValue;
use datafusion::common::{DataFusionError, Result as DataFusionResult};
use datafusion::logical_expr::ColumnarValue;
use datafusion::physical_expr::PhysicalExpr;
use std::{
Expand Down Expand Up @@ -142,40 +138,6 @@ fn relabel_array(array: ArrayRef, target_type: &DataType) -> ArrayRef {
}
}

/// Casts a Timestamp(Microsecond) array to Timestamp(Millisecond) by dividing values by 1000.
/// Preserves the timezone from the target type.
fn cast_timestamp_micros_to_millis_array(
array: &ArrayRef,
target_tz: Option<Arc<str>>,
) -> ArrayRef {
let micros_array = array
.as_any()
.downcast_ref::<TimestampMicrosecondArray>()
.expect("Expected TimestampMicrosecondArray");

let millis_values: TimestampMillisecondArray =
arrow::compute::kernels::arity::unary(micros_array, |v| v / 1000);

// Apply timezone if present
let result = if let Some(tz) = target_tz {
millis_values.with_timezone(tz)
} else {
millis_values
};

Arc::new(result)
}

/// Casts a Timestamp(Microsecond) scalar to Timestamp(Millisecond) by dividing the value by 1000.
/// Preserves the timezone from the target type.
fn cast_timestamp_micros_to_millis_scalar(
opt_val: Option<i64>,
target_tz: Option<Arc<str>>,
) -> ScalarValue {
let new_val = opt_val.map(|v| v / 1000);
ScalarValue::TimestampMillisecond(new_val, target_tz)
}

#[derive(Debug, Clone, Eq)]
pub struct CometCastColumnExpr {
/// The physical expression producing the value to cast.
Expand Down Expand Up @@ -214,20 +176,41 @@ impl Hash for CometCastColumnExpr {
}

impl CometCastColumnExpr {
/// Create a new [`CometCastColumnExpr`].
pub fn new(
/// Try to create a new [`CometCastColumnExpr`].
pub fn try_new(
expr: Arc<dyn PhysicalExpr>,
physical_field: FieldRef,
target_field: FieldRef,
cast_options: Option<CastOptions<'static>>,
) -> Self {
Self {
) -> DataFusionResult<Self> {
let physical_type = physical_field.data_type();
let target_type = target_field.data_type();
// `target_field` is the Spark logical field, while `physical_field` comes from the
// Parquet or Iceberg file. Comet represents Spark's TimestampType and TimestampNTZType
// as Arrow microseconds, and Spark maps both TIMESTAMP_MICROS and TIMESTAMP_MILLIS files
// to those logical types. A millisecond target is therefore invalid at this read-adapter
// boundary:
// https://github.com/apache/spark/blob/v4.2.0/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetSchemaConverter.scala#L318-L324
if matches!(
(physical_type, target_type),
(
DataType::Timestamp(TimeUnit::Microsecond, _),
DataType::Timestamp(TimeUnit::Millisecond, _)
)
) {
return Err(DataFusionError::Plan(format!(
"Cannot adapt Spark timestamp field '{}' from {physical_type} to {target_type}: Spark read schemas represent logical timestamps in microseconds",
physical_field.name()
)));
}

Ok(Self {
expr,
input_physical_field: physical_field,
target_field,
cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS),
parquet_options: None,
}
})
}

/// Set Spark parquet options to enable complex nested type conversions.
Expand Down Expand Up @@ -271,23 +254,7 @@ impl PhysicalExpr for CometCastColumnExpr {
let input_physical_field = self.input_physical_field.data_type();
let target_field = self.target_field.data_type();

// Handle specific type conversions with custom casts
match (input_physical_field, target_field) {
// Timestamp(Microsecond) -> Timestamp(Millisecond)
(
DataType::Timestamp(TimeUnit::Microsecond, _),
DataType::Timestamp(TimeUnit::Millisecond, target_tz),
) => match value {
ColumnarValue::Array(array) => {
let casted = cast_timestamp_micros_to_millis_array(&array, target_tz.clone());
Ok(ColumnarValue::Array(casted))
}
ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(opt_val, _)) => {
let casted = cast_timestamp_micros_to_millis_scalar(opt_val, target_tz.clone());
Ok(ColumnarValue::Scalar(casted))
}
_ => Ok(value),
},
// Nested types that differ only in field names (e.g., List element named
// "item" vs "element", or Map entries named "key_value" vs "entries").
// Re-label the array so the DataType metadata matches the logical schema.
Expand Down Expand Up @@ -329,12 +296,12 @@ impl PhysicalExpr for CometCastColumnExpr {
) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
assert_eq!(children.len(), 1);
let child = children.pop().expect("CastColumnExpr child");
let mut new_expr = Self::new(
let mut new_expr = Self::try_new(
child,
Arc::clone(&self.input_physical_field),
Arc::clone(&self.target_field),
Some(self.cast_options.clone()),
);
)?;
if let Some(opts) = &self.parquet_options {
new_expr = new_expr.with_parquet_options(opts.clone());
}
Expand All @@ -354,154 +321,27 @@ mod tests {
use datafusion::physical_expr::expressions::Column;

#[test]
fn test_cast_timestamp_micros_to_millis_array() {
// Create a TimestampMicrosecond array with some values
let micros_array: TimestampMicrosecondArray = vec![
Some(1_000_000), // 1 second in micros
Some(2_500_000), // 2.5 seconds in micros
None, // null value
Some(0), // zero
Some(-1_000_000), // negative value (before epoch)
]
.into();
let array_ref: ArrayRef = Arc::new(micros_array);

// Cast without timezone
let result = cast_timestamp_micros_to_millis_array(&array_ref, None);
let millis_array = result
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.expect("Expected TimestampMillisecondArray");

assert_eq!(millis_array.len(), 5);
assert_eq!(millis_array.value(0), 1000); // 1_000_000 / 1000
assert_eq!(millis_array.value(1), 2500); // 2_500_000 / 1000
assert!(millis_array.is_null(2));
assert_eq!(millis_array.value(3), 0);
assert_eq!(millis_array.value(4), -1000); // -1_000_000 / 1000
}

#[test]
fn test_cast_timestamp_micros_to_millis_array_with_timezone() {
let micros_array: TimestampMicrosecondArray = vec![Some(1_000_000), Some(2_000_000)].into();
let array_ref: ArrayRef = Arc::new(micros_array);

let target_tz: Option<Arc<str>> = Some(Arc::from("UTC"));
let result = cast_timestamp_micros_to_millis_array(&array_ref, target_tz);
let millis_array = result
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.expect("Expected TimestampMillisecondArray");

assert_eq!(millis_array.value(0), 1000);
assert_eq!(millis_array.value(1), 2000);
// Verify timezone is preserved
assert_eq!(
result.data_type(),
&DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("UTC")))
);
}

#[test]
fn test_cast_timestamp_micros_to_millis_scalar() {
// Test with a value
let result = cast_timestamp_micros_to_millis_scalar(Some(1_500_000), None);
assert_eq!(result, ScalarValue::TimestampMillisecond(Some(1500), None));

// Test with null
let null_result = cast_timestamp_micros_to_millis_scalar(None, None);
assert_eq!(null_result, ScalarValue::TimestampMillisecond(None, None));

// Test with timezone
let target_tz: Option<Arc<str>> = Some(Arc::from("UTC"));
let tz_result = cast_timestamp_micros_to_millis_scalar(Some(2_000_000), target_tz.clone());
assert_eq!(
tz_result,
ScalarValue::TimestampMillisecond(Some(2000), target_tz)
);
}

#[test]
fn test_comet_cast_column_expr_evaluate_micros_to_millis_array() {
// Create input schema with TimestampMicrosecond column
let input_field = Arc::new(Field::new(
"ts",
DataType::Timestamp(TimeUnit::Microsecond, None),
true,
));
let schema = Schema::new(vec![Arc::clone(&input_field)]);

// Create target field with TimestampMillisecond
let target_field = Arc::new(Field::new(
"ts",
DataType::Timestamp(TimeUnit::Millisecond, None),
true,
));

// Create a column expression
let col_expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("ts", 0));

// Create the CometCastColumnExpr
let cast_expr = CometCastColumnExpr::new(col_expr, input_field, target_field, None);

// Create a record batch with TimestampMicrosecond data
let micros_array: TimestampMicrosecondArray =
vec![Some(1_000_000), Some(2_000_000), None].into();
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(micros_array)]).unwrap();

// Evaluate
let result = cast_expr.evaluate(&batch).unwrap();

match result {
ColumnarValue::Array(arr) => {
let millis_array = arr
.as_any()
.downcast_ref::<TimestampMillisecondArray>()
.expect("Expected TimestampMillisecondArray");
assert_eq!(millis_array.value(0), 1000);
assert_eq!(millis_array.value(1), 2000);
assert!(millis_array.is_null(2));
}
_ => panic!("Expected Array result"),
}
}

#[test]
fn test_comet_cast_column_expr_evaluate_micros_to_millis_scalar() {
// Create input schema with TimestampMicrosecond column
let input_field = Arc::new(Field::new(
"ts",
DataType::Timestamp(TimeUnit::Microsecond, None),
true,
));
let schema = Schema::new(vec![Arc::clone(&input_field)]);

// Create target field with TimestampMillisecond
let target_field = Arc::new(Field::new(
"ts",
DataType::Timestamp(TimeUnit::Millisecond, None),
true,
));

// Create a literal expression that returns a scalar
let scalar = ScalarValue::TimestampMicrosecond(Some(1_500_000), None);
let literal_expr: Arc<dyn PhysicalExpr> =
Arc::new(datafusion::physical_expr::expressions::Literal::new(scalar));

// Create the CometCastColumnExpr
let cast_expr = CometCastColumnExpr::new(literal_expr, input_field, target_field, None);

// Create an empty batch (scalar doesn't need data)
let batch = RecordBatch::new_empty(Arc::new(schema));

// Evaluate
let result = cast_expr.evaluate(&batch).unwrap();

match result {
ColumnarValue::Scalar(s) => {
assert_eq!(s, ScalarValue::TimestampMillisecond(Some(1500), None));
}
_ => panic!("Expected Scalar result"),
fn test_rejects_millisecond_logical_timestamp() {
for timezone in [None, Some(Arc::from("UTC"))] {
let input_field = Arc::new(Field::new(
"ts",
DataType::Timestamp(TimeUnit::Microsecond, timezone.clone()),
true,
));
let target_field = Arc::new(Field::new(
"ts",
DataType::Timestamp(TimeUnit::Millisecond, timezone),
true,
));
let expr: Arc<dyn PhysicalExpr> = Arc::new(Column::new("ts", 0));

let err = CometCastColumnExpr::try_new(expr, input_field, target_field, None)
.expect_err("millisecond logical timestamp must be rejected during planning");
assert!(matches!(
err,
DataFusionError::Plan(message)
if message.contains("Spark read schemas represent logical timestamps in microseconds")
));
}
}

Expand Down
8 changes: 4 additions & 4 deletions native/core/src/parquet/schema_adapter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -603,12 +603,12 @@ impl SparkPhysicalExprAdapter {
}

let cast_expr: Arc<dyn PhysicalExpr> = Arc::new(
CometCastColumnExpr::new(
CometCastColumnExpr::try_new(
remapped,
Arc::clone(physical_field),
Arc::clone(logical_field),
None,
)
)?
.with_parquet_options(self.parquet_options.clone()),
);
return Ok(Transformed::yes(cast_expr));
Expand Down Expand Up @@ -892,12 +892,12 @@ impl SparkPhysicalExprAdapter {
| (DataType::Timestamp(_, _), DataType::Int64)
) {
let comet_cast: Arc<dyn PhysicalExpr> = Arc::new(
CometCastColumnExpr::new(
CometCastColumnExpr::try_new(
child,
input_field,
Arc::clone(cast.target_field()),
None,
)
)?
.with_parquet_options(self.parquet_options.clone()),
);
return Ok(Transformed::yes(comet_cast));
Expand Down
Loading
Loading