Skip to content

Commit 6e5344a

Browse files
xinlifoobarxinlifoobar
authored
Extract parquet statistics from timestamps with timezones (#10766)
* Fix incorrect statistics read for timestamp columns in parquet
1 parent 888504a commit 6e5344a

3 files changed

Lines changed: 615 additions & 99 deletions

File tree

‎datafusion/core/src/datasource/physical_plan/parquet/statistics.rs‎

Lines changed: 263 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
2020
// TODO: potentially move this to arrow-rs: https://github.com/apache/arrow-rs/issues/4328
2121

22-
use arrow::{array::ArrayRef, datatypes::DataType};
22+
use arrow::{array::ArrayRef, datatypes::DataType, datatypes::TimeUnit};
2323
use arrow_array::{new_empty_array, new_null_array, UInt64Array};
2424
use arrow_schema::{Field, FieldRef, Schema};
2525
use datafusion_common::{
@@ -112,6 +112,26 @@ macro_rules! get_statistic {
112112
Some(DataType::UInt64) => {
113113
Some(ScalarValue::UInt64(Some((*s.$func()) as u64)))
114114
}
115+
Some(DataType::Timestamp(unit, timezone)) => {
116+
Some(match unit {
117+
TimeUnit::Second => ScalarValue::TimestampSecond(
118+
Some(*s.$func()),
119+
timezone.clone(),
120+
),
121+
TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(
122+
Some(*s.$func()),
123+
timezone.clone(),
124+
),
125+
TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(
126+
Some(*s.$func()),
127+
timezone.clone(),
128+
),
129+
TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(
130+
Some(*s.$func()),
131+
timezone.clone(),
132+
),
133+
})
134+
}
115135
_ => Some(ScalarValue::Int64(Some(*s.$func()))),
116136
}
117137
}
@@ -395,7 +415,8 @@ mod test {
395415
use arrow_array::{
396416
new_null_array, Array, BinaryArray, BooleanArray, Date32Array, Date64Array,
397417
Decimal128Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array,
398-
Int8Array, RecordBatch, StringArray, StructArray, TimestampNanosecondArray,
418+
Int8Array, RecordBatch, StringArray, StructArray, TimestampMicrosecondArray,
419+
TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
399420
};
400421
use arrow_schema::{Field, SchemaRef};
401422
use bytes::Bytes;
@@ -536,28 +557,209 @@ mod test {
536557
}
537558

538559
#[test]
539-
#[should_panic(
540-
expected = "Inconsistent types in ScalarValue::iter_to_array. Expected Int64, got TimestampNanosecond(NULL, None)"
541-
)]
542-
// Due to https://github.com/apache/datafusion/issues/8295
543560
fn roundtrip_timestamp() {
544561
Test {
545-
input: timestamp_array([
546-
// row group 1
547-
Some(1),
548-
None,
549-
Some(3),
550-
// row group 2
551-
Some(9),
552-
Some(5),
562+
input: timestamp_seconds_array(
563+
[
564+
// row group 1
565+
Some(1),
566+
None,
567+
Some(3),
568+
// row group 2
569+
Some(9),
570+
Some(5),
571+
None,
572+
// row group 3
573+
None,
574+
None,
575+
None,
576+
],
553577
None,
554-
// row group 3
578+
),
579+
expected_min: timestamp_seconds_array([Some(1), Some(5), None], None),
580+
expected_max: timestamp_seconds_array([Some(3), Some(9), None], None),
581+
}
582+
.run();
583+
584+
Test {
585+
input: timestamp_milliseconds_array(
586+
[
587+
// row group 1
588+
Some(1),
589+
None,
590+
Some(3),
591+
// row group 2
592+
Some(9),
593+
Some(5),
594+
None,
595+
// row group 3
596+
None,
597+
None,
598+
None,
599+
],
555600
None,
601+
),
602+
expected_min: timestamp_milliseconds_array([Some(1), Some(5), None], None),
603+
expected_max: timestamp_milliseconds_array([Some(3), Some(9), None], None),
604+
}
605+
.run();
606+
607+
Test {
608+
input: timestamp_microseconds_array(
609+
[
610+
// row group 1
611+
Some(1),
612+
None,
613+
Some(3),
614+
// row group 2
615+
Some(9),
616+
Some(5),
617+
None,
618+
// row group 3
619+
None,
620+
None,
621+
None,
622+
],
556623
None,
624+
),
625+
expected_min: timestamp_microseconds_array([Some(1), Some(5), None], None),
626+
expected_max: timestamp_microseconds_array([Some(3), Some(9), None], None),
627+
}
628+
.run();
629+
630+
Test {
631+
input: timestamp_nanoseconds_array(
632+
[
633+
// row group 1
634+
Some(1),
635+
None,
636+
Some(3),
637+
// row group 2
638+
Some(9),
639+
Some(5),
640+
None,
641+
// row group 3
642+
None,
643+
None,
644+
None,
645+
],
557646
None,
558-
]),
559-
expected_min: timestamp_array([Some(1), Some(5), None]),
560-
expected_max: timestamp_array([Some(3), Some(9), None]),
647+
),
648+
expected_min: timestamp_nanoseconds_array([Some(1), Some(5), None], None),
649+
expected_max: timestamp_nanoseconds_array([Some(3), Some(9), None], None),
650+
}
651+
.run()
652+
}
653+
654+
#[test]
655+
fn roundtrip_timestamp_timezoned() {
656+
Test {
657+
input: timestamp_seconds_array(
658+
[
659+
// row group 1
660+
Some(1),
661+
None,
662+
Some(3),
663+
// row group 2
664+
Some(9),
665+
Some(5),
666+
None,
667+
// row group 3
668+
None,
669+
None,
670+
None,
671+
],
672+
Some("UTC"),
673+
),
674+
expected_min: timestamp_seconds_array([Some(1), Some(5), None], Some("UTC")),
675+
expected_max: timestamp_seconds_array([Some(3), Some(9), None], Some("UTC")),
676+
}
677+
.run();
678+
679+
Test {
680+
input: timestamp_milliseconds_array(
681+
[
682+
// row group 1
683+
Some(1),
684+
None,
685+
Some(3),
686+
// row group 2
687+
Some(9),
688+
Some(5),
689+
None,
690+
// row group 3
691+
None,
692+
None,
693+
None,
694+
],
695+
Some("UTC"),
696+
),
697+
expected_min: timestamp_milliseconds_array(
698+
[Some(1), Some(5), None],
699+
Some("UTC"),
700+
),
701+
expected_max: timestamp_milliseconds_array(
702+
[Some(3), Some(9), None],
703+
Some("UTC"),
704+
),
705+
}
706+
.run();
707+
708+
Test {
709+
input: timestamp_microseconds_array(
710+
[
711+
// row group 1
712+
Some(1),
713+
None,
714+
Some(3),
715+
// row group 2
716+
Some(9),
717+
Some(5),
718+
None,
719+
// row group 3
720+
None,
721+
None,
722+
None,
723+
],
724+
Some("UTC"),
725+
),
726+
expected_min: timestamp_microseconds_array(
727+
[Some(1), Some(5), None],
728+
Some("UTC"),
729+
),
730+
expected_max: timestamp_microseconds_array(
731+
[Some(3), Some(9), None],
732+
Some("UTC"),
733+
),
734+
}
735+
.run();
736+
737+
Test {
738+
input: timestamp_nanoseconds_array(
739+
[
740+
// row group 1
741+
Some(1),
742+
None,
743+
Some(3),
744+
// row group 2
745+
Some(9),
746+
Some(5),
747+
None,
748+
// row group 3
749+
None,
750+
None,
751+
None,
752+
],
753+
Some("UTC"),
754+
),
755+
expected_min: timestamp_nanoseconds_array(
756+
[Some(1), Some(5), None],
757+
Some("UTC"),
758+
),
759+
expected_max: timestamp_nanoseconds_array(
760+
[Some(3), Some(9), None],
761+
Some("UTC"),
762+
),
561763
}
562764
.run()
563765
}
@@ -914,8 +1116,8 @@ mod test {
9141116
// File has no min/max for timestamp_col
9151117
.with_column(ExpectedColumn {
9161118
name: "timestamp_col",
917-
expected_min: timestamp_array([None]),
918-
expected_max: timestamp_array([None]),
1119+
expected_min: timestamp_nanoseconds_array([None], None),
1120+
expected_max: timestamp_nanoseconds_array([None], None),
9191121
})
9201122
.with_column(ExpectedColumn {
9211123
name: "year",
@@ -1135,9 +1337,48 @@ mod test {
11351337
Arc::new(array)
11361338
}
11371339

1138-
fn timestamp_array(input: impl IntoIterator<Item = Option<i64>>) -> ArrayRef {
1340+
fn timestamp_seconds_array(
1341+
input: impl IntoIterator<Item = Option<i64>>,
1342+
timzezone: Option<&str>,
1343+
) -> ArrayRef {
1344+
let array: TimestampSecondArray = input.into_iter().collect();
1345+
match timzezone {
1346+
Some(tz) => Arc::new(array.with_timezone(tz)),
1347+
None => Arc::new(array),
1348+
}
1349+
}
1350+
1351+
fn timestamp_milliseconds_array(
1352+
input: impl IntoIterator<Item = Option<i64>>,
1353+
timzezone: Option<&str>,
1354+
) -> ArrayRef {
1355+
let array: TimestampMillisecondArray = input.into_iter().collect();
1356+
match timzezone {
1357+
Some(tz) => Arc::new(array.with_timezone(tz)),
1358+
None => Arc::new(array),
1359+
}
1360+
}
1361+
1362+
fn timestamp_microseconds_array(
1363+
input: impl IntoIterator<Item = Option<i64>>,
1364+
timzezone: Option<&str>,
1365+
) -> ArrayRef {
1366+
let array: TimestampMicrosecondArray = input.into_iter().collect();
1367+
match timzezone {
1368+
Some(tz) => Arc::new(array.with_timezone(tz)),
1369+
None => Arc::new(array),
1370+
}
1371+
}
1372+
1373+
fn timestamp_nanoseconds_array(
1374+
input: impl IntoIterator<Item = Option<i64>>,
1375+
timzezone: Option<&str>,
1376+
) -> ArrayRef {
11391377
let array: TimestampNanosecondArray = input.into_iter().collect();
1140-
Arc::new(array)
1378+
match timzezone {
1379+
Some(tz) => Arc::new(array.with_timezone(tz)),
1380+
None => Arc::new(array),
1381+
}
11411382
}
11421383

11431384
fn utf8_array<'a>(input: impl IntoIterator<Item = Option<&'a str>>) -> ArrayRef {

0 commit comments

Comments
 (0)