ump45nose commented on code in PR #7089:
URL: https://github.com/apache/datafusion-comet/pull/7089#discussion_r4332447504
##########
native/core/src/parquet/schema_adapter.rs:
##########
@@ -2198,6 +2264,366 @@ pub(crate) mod test {
Ok(())
}
+ /// `TIMESTAMP(NANOS)` read as `timestamp` or `timestamp_ntz`: Spark reads
nanosecond
+ /// timestamps only as `bigint` (`nanosAsLong`), so the cast to
microseconds is rejected
+ /// whether or not the file column is UTC-adjusted. See #7034.
+ #[tokio::test]
+ async fn parquet_nanos_timestamp_read_as_timestamp_errors() -> Result<(),
DataFusionError> {
+ for file_tz in [Some("UTC"), None] {
+ for read_tz in [Some("UTC"), None] {
+ let values = Arc::new(
+
TimestampNanosecondArray::from(vec![1_600_000_000_123_456_789i64])
+ .with_timezone_opt(file_tz),
+ ) as Arc<dyn arrow::array::Array>;
+ let msg = assert_rejected_conversion(
+ Field::new(
+ "t",
+ DataType::Timestamp(TimeUnit::Nanosecond,
file_tz.map(Into::into)),
+ false,
+ ),
+ values,
+ DataType::Timestamp(TimeUnit::Microsecond,
read_tz.map(Into::into)),
+ )
+ .await?;
+ assert!(
+ msg.contains("Column: [[t]]")
+ && msg.contains("Found: INT64")
+ && msg.contains("Expected: timestamp"),
+ "file tz {file_tz:?}, read tz {read_tz:?}: unexpected
error: {msg}"
+ );
+ }
+ }
+ Ok(())
+ }
+
+ /// The `TIMESTAMP(NANOS)` rejection keys on the footer check's marker,
not on the Arrow
+ /// type: a marked leaf is rejected at any depth and with or without a
timezone, and an
+ /// unmarked nanosecond leaf, which may be INT96, is accepted. Reading
nanoseconds as `bigint`
+ /// stays allowed.
+ #[test]
+ fn nanos_timestamp_rejection_keys_on_the_parquet_type() -> Result<(),
DataFusionError> {
+ let options = default_options();
+ let nanos_field = |name: &str, tz: Option<&str>, marked: bool| {
+ let field = Field::new(
+ name,
+ DataType::Timestamp(TimeUnit::Nanosecond, tz.map(Into::into)),
+ true,
+ );
+ if marked {
+ field.with_metadata(std::collections::HashMap::from([(
+ NANOS_TIMESTAMP_KEY.to_string(),
+ "true".to_string(),
+ )]))
+ } else {
+ field
+ }
+ };
+ let micros_utc = DataType::Timestamp(TimeUnit::Microsecond,
Some("UTC".into()));
+ let micros_ntz = DataType::Timestamp(TimeUnit::Microsecond, None);
+ for tz in [None, Some("UTC")] {
+ for marked in [true, false] {
+ let leaf = || Arc::new(nanos_field("t", tz, marked));
+ let element = || Arc::new(nanos_field("element", tz, marked));
+ let map_of = |value: Field| {
+ DataType::Map(
+ Arc::new(Field::new(
+ "key_value",
+ DataType::Struct(Fields::from(vec![
+ Field::new("key", DataType::Utf8, false),
+ value,
+ ])),
+ false,
+ )),
+ false,
+ )
+ };
+ for (physical, target) in [
+ (leaf().as_ref().clone(), micros_utc.clone()),
+ (leaf().as_ref().clone(), micros_ntz.clone()),
+ (
+ Field::new("a", DataType::List(element()), true),
+ list_type(micros_utc.clone()),
+ ),
+ (
+ Field::new("a",
DataType::Struct(Fields::from(vec![leaf()])), true),
+ struct_type(vec![("t", micros_utc.clone())]),
+ ),
+ (
+ Field::new(
+ "a",
+ DataType::List(Arc::new(Field::new(
+ "element",
+ DataType::Struct(Fields::from(vec![leaf()])),
+ true,
+ ))),
+ true,
+ ),
+ list_type(struct_type(vec![("t",
micros_utc.clone())])),
+ ),
+ (
+ Field::new("a", map_of(nanos_field("value", tz,
marked)), true),
+ map_of(Field::new("value", micros_utc.clone(), true)),
+ ),
+ ] {
+ // Reading a UTC-adjusted timestamp as `timestamp_ntz` is
rejected anyway
+ // with the default (Spark 3.x) options; see SPARK-36182.
+ let expected = marked || (tz.is_some() && target ==
micros_ntz);
+ let check = check_field_conversion(&physical, &target,
"a", &options)?;
+ assert_eq!(
+ matches!(check, ConversionCheck::RejectOnNonEmpty { ..
}),
+ expected,
+ "{physical:?} (marked: {marked}) read as {target}"
+ );
+ }
+ let check = check_field_conversion(
+ &leaf().as_ref().clone(),
+ &DataType::Int64,
+ "t",
+ &options,
+ )?;
+ assert!(matches!(check, ConversionCheck::Accept), "{tz:?} as
bigint");
+ }
+ }
+ Ok(())
+ }
+
+ /// An INT96 array in the legacy `bag` list layout reaches the conversion
check as a
+ /// timezone-free nanosecond leaf below a list (#7029), the same Arrow
type as a
+ /// `TIMESTAMP(NANOS)` leaf, so the `TIMESTAMP(NANOS)` rejection must
leave it alone. Write
+ /// such a file and read it through the scan.
+ #[tokio::test]
+ async fn legacy_int96_array_still_reads_as_timestamp() -> Result<(),
DataFusionError> {
+ use parquet::data_type::{Int96, Int96Type};
+ use parquet::file::writer::SerializedFileWriter;
+ use parquet::schema::parser::parse_message_type;
+
+ let schema = Arc::new(parse_message_type(
+ "message spark_schema {
+ OPTIONAL GROUP a (LIST) {
+ REPEATED GROUP bag { OPTIONAL INT96 array; }
+ }
+ }",
+ )?);
+ // 1970-01-02 00:00:00.001 and 1970-01-03 00:00:00: INT96 holds the
nanoseconds of the
+ // day (low word first) and then the Julian day.
+ const UNIX_EPOCH_JULIAN_DAY: u32 = 2_440_588;
+ let mut first = Int96::new();
+ first.set_data(1_000_000, 0, UNIX_EPOCH_JULIAN_DAY + 1);
+ let mut second = Int96::new();
+ second.set_data(0, 0, UNIX_EPOCH_JULIAN_DAY + 2);
+
+ let filename = get_temp_filename().to_str().unwrap().to_string();
+ let mut writer =
+ SerializedFileWriter::new(File::create(&filename)?, schema,
Default::default())?;
+ let mut row_group = writer.next_row_group()?;
+ let mut a = row_group.next_column()?.unwrap();
+ a.typed::<Int96Type>()
+ .write_batch(&[first, second], Some(&[3, 3]), Some(&[0, 1]))?;
+ a.close()?;
+ row_group.close()?;
+ writer.close()?;
+
+ let micros_utc = DataType::Timestamp(TimeUnit::Microsecond,
Some("UTC".into()));
+ let required = Arc::new(Schema::new(vec![Field::new(
+ "a",
+ list_type(micros_utc),
+ true,
+ )]));
+ let mut stream = scan_parquet_file(filename, required,
default_options(), None, false)?;
+ let batch = stream.next().await.unwrap()?;
+ assert!(stream.next().await.is_none());
+
+ let list = batch.column(0).as_list::<i32>();
+ assert_eq!(list.len(), 1);
+ let values = list.value(0);
+ let expected = TimestampMicrosecondArray::from(vec![86_400_001_000i64,
172_800_000_000])
+ .with_timezone("UTC");
+ assert_eq!(
+ values
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .expect("timestamp(us, UTC) values"),
+ &expected
+ );
+ Ok(())
+ }
+
+ /// Write `id INT32 = 1` and `t INT96 = 1970-01-02 00:00:00.001` with the
low-level Parquet
+ /// writer, adding `arrow_schema` to the footer as the `ARROW:schema` hint
when given.
+ fn write_int96_file(arrow_schema: Option<&Schema>) -> Result<String,
DataFusionError> {
+ use parquet::arrow::{encode_arrow_schema, ARROW_SCHEMA_META_KEY};
+ use parquet::data_type::{Int32Type as ParquetInt32, Int96, Int96Type};
+ use parquet::file::metadata::KeyValue;
+ use parquet::file::properties::WriterProperties;
+ use parquet::file::writer::SerializedFileWriter;
+ use parquet::schema::parser::parse_message_type;
+
+ let schema = Arc::new(parse_message_type(
+ "message spark_schema { REQUIRED INT32 id; OPTIONAL INT96 t; }",
+ )?);
+ let properties = WriterProperties::builder()
+ .set_key_value_metadata(arrow_schema.map(|arrow_schema| {
+ vec![KeyValue::new(
+ ARROW_SCHEMA_META_KEY.to_string(),
+ encode_arrow_schema(arrow_schema),
+ )]
+ }))
+ .build();
+ let mut t = Int96::new();
+ t.set_data(1_000_000, 0, 2_440_589);
+
+ let filename = get_temp_filename().to_str().unwrap().to_string();
+ let mut writer =
+ SerializedFileWriter::new(File::create(&filename)?, schema,
Arc::new(properties))?;
+ let mut row_group = writer.next_row_group()?;
+ let mut id = row_group.next_column()?.unwrap();
+ id.typed::<ParquetInt32>().write_batch(&[1], None, None)?;
+ id.close()?;
+ let mut column = row_group.next_column()?.unwrap();
+ column
+ .typed::<Int96Type>()
+ .write_batch(&[t], Some(&[1]), None)?;
+ column.close()?;
+ row_group.close()?;
+ writer.close()?;
+ Ok(filename)
+ }
+
+ /// Read `id` and `t` of [`write_int96_file`] as `int` and `timestamp`
with the filter
+ /// `id > 0`, with or without row filters, and check that `t` comes back.
+ async fn assert_int96_reads_as_timestamp(filename: &str) -> Result<(),
DataFusionError> {
+ let micros_utc = DataType::Timestamp(TimeUnit::Microsecond,
Some("UTC".into()));
+ let required = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("t", micros_utc, true),
+ ]));
+ for row_filters in [false, true] {
+ let id_positive: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("id", 0)),
+ Operator::Gt,
+ Arc::new(Literal::new(ScalarValue::Int32(Some(0)))),
+ ));
+ let mut stream = scan_parquet_file(
+ filename.to_string(),
+ Arc::clone(&required),
+ default_options(),
+ Some(id_positive),
+ row_filters,
+ )?;
+ let batch = stream.next().await.unwrap()?;
+ assert!(stream.next().await.is_none());
+ assert_eq!(
+ batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .expect("timestamp(us, UTC) values"),
+
&TimestampMicrosecondArray::from(vec![86_400_001_000i64]).with_timezone("UTC"),
+ "row filters: {row_filters}"
+ );
+ }
+ Ok(())
+ }
+
+ /// A plain INT96 column reads as `timestamp`, also with row filters,
where the footer
+ /// check sees it before INT96 coercion, as a timezone-free nanosecond
timestamp.
+ #[tokio::test]
+ async fn int96_reads_as_timestamp_with_row_filters() -> Result<(),
DataFusionError> {
+ assert_int96_reads_as_timestamp(&write_int96_file(None)?).await
+ }
+
+ /// An INT96 column whose `ARROW:schema` hint says `timestamp('ns',
tz='UTC')`, as PyArrow
+ /// writes with `use_deprecated_int96_timestamps=True`, keeps that
timezone through INT96
+ /// coercion, and still reads as `timestamp`.
+ #[tokio::test]
+ async fn int96_with_arrow_timezone_hint_reads_as_timestamp() -> Result<(),
DataFusionError> {
+ let hint = Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new(
+ "t",
+ DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
+ true,
+ ),
+ ]);
+ assert_int96_reads_as_timestamp(&write_int96_file(Some(&hint))?).await
+ }
+
+ /// `TIMESTAMP(NANOS)` is rejected below a list too, with or without row
filters and
+ /// whether or not it is UTC-adjusted, since the rule no longer has to
leave timezone-free
+ /// nanosecond leaves below a list to INT96.
+ #[tokio::test]
+ async fn nanos_timestamp_array_read_as_timestamp_errors() -> Result<(),
DataFusionError> {
Review Comment:
Agreed, it only added the column path. Removed in a7eb8e9; the `[[a, list,
element]]` path assertion moved into
`parquet_nanos_timestamp_read_as_timestamp_errors`. The slot now holds the
ListView leaf-count scan test, which the JVM side can't target as directly.
--
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]
