Implement Parquet GH-583 INT96 timestamp ColumnOrder

parquet/src/basic.rs

@@ -987,6 +987,11 @@ pub enum SortOrder {
     UNDEFINED,
     /// Use IEEE 754 total order.
     TOTAL_ORDER,
+    /// Use INT96 timestamp order (see [parquet-format/#584] and the [Thrift spec]).
+    ///
+    /// [parquet-format/#584]: https://github.com/apache/parquet-format/pull/584
+    /// [Thrift spec]: https://github.com/apache/parquet-format/blob/2076361bb64e2de9ca6a8d06eda025a6fa4e9df6/src/main/thrift/parquet.thrift#L1230-L1233
+    INT96_TIMESTAMP,
 }
 
 impl SortOrder {
@@ -1009,6 +1014,8 @@ pub enum ColumnOrder {
     TYPE_DEFINED_ORDER(SortOrder),
     /// Column ordering to use for floating point types.
     IEEE_754_TOTAL_ORDER,
+    /// Column ordering to use for INT96 types.
+    INT96_TIMESTAMP_ORDER,
     // The following are not defined in the Parquet spec and should always be last.
     /// Undefined column order, means legacy behaviour before parquet-format 2.4.0.
     /// Sort order is always SIGNED.
@@ -1043,6 +1050,8 @@ impl ColumnOrder {
             || matches!(physical_type, Type::FLOAT | Type::DOUBLE)
         {
             ColumnOrder::IEEE_754_TOTAL_ORDER
+        } else if matches!(physical_type, Type::INT96) {
+            ColumnOrder::INT96_TIMESTAMP_ORDER
         } else {
             let sort_order =
                 Self::get_sort_order_for_type(logical_type, converted_type, physical_type, true);
@@ -1161,7 +1170,13 @@ impl ColumnOrder {
             // Order: false, true
             Type::BOOLEAN => SortOrder::UNSIGNED,
             Type::INT32 | Type::INT64 => SortOrder::SIGNED,
-            Type::INT96 => SortOrder::UNDEFINED,
+            Type::INT96 => {
+                if is_type_defined {
+                    SortOrder::UNDEFINED
+                } else {
+                    SortOrder::INT96_TIMESTAMP
+                }
+            }
             // Notes to remember when comparing float/double values:
             // If legacy TYPE_DEFINED_ORDER is specified:
             //   If the min is a NaN, it should be ignored.
@@ -1190,6 +1205,7 @@ impl ColumnOrder {
         match *self {
             ColumnOrder::TYPE_DEFINED_ORDER(order) => order,
             ColumnOrder::IEEE_754_TOTAL_ORDER => SortOrder::TOTAL_ORDER,
+            ColumnOrder::INT96_TIMESTAMP_ORDER => SortOrder::INT96_TIMESTAMP,
             ColumnOrder::UNDEFINED => SortOrder::SIGNED,
             ColumnOrder::UNKNOWN => SortOrder::UNDEFINED,
         }
@@ -1212,6 +1228,10 @@ impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for ColumnOrder {
                 prot.skip_empty_struct()?;
                 Self::IEEE_754_TOTAL_ORDER
             }
+            3 => {
+                prot.skip_empty_struct()?;
+                Self::INT96_TIMESTAMP_ORDER
+            }
             _ => {
                 prot.skip(field_ident.field_type)?;
                 Self::UNKNOWN
@@ -1240,6 +1260,10 @@ impl WriteThrift for ColumnOrder {
                 writer.write_field_begin(FieldType::Struct, 2, 0)?;
                 writer.write_struct_end()?;
             }
+            Self::INT96_TIMESTAMP_ORDER => {
+                writer.write_field_begin(FieldType::Struct, 3, 0)?;
+                writer.write_struct_end()?;
+            }
             _ => return Err(general_err!("Attempt to write undefined ColumnOrder")),
         }
         // write end of struct for this union
@@ -2040,6 +2064,7 @@ mod tests {
         assert_eq!(SortOrder::UNSIGNED.to_string(), "UNSIGNED");
         assert_eq!(SortOrder::UNDEFINED.to_string(), "UNDEFINED");
         assert_eq!(SortOrder::TOTAL_ORDER.to_string(), "TOTAL_ORDER");
+        assert_eq!(SortOrder::INT96_TIMESTAMP.to_string(), "INT96_TIMESTAMP");
     }
 
     #[test]
@@ -2060,6 +2085,10 @@ mod tests {
             ColumnOrder::IEEE_754_TOTAL_ORDER.to_string(),
             "IEEE_754_TOTAL_ORDER"
         );
+        assert_eq!(
+            ColumnOrder::INT96_TIMESTAMP_ORDER.to_string(),
+            "INT96_TIMESTAMP_ORDER"
+        );
         assert_eq!(ColumnOrder::UNDEFINED.to_string(), "UNDEFINED");
     }
 
@@ -2206,6 +2235,10 @@ mod tests {
             ColumnOrder::get_default_sort_order(Type::INT96, true),
             SortOrder::UNDEFINED
         );
+        assert_eq!(
+            ColumnOrder::get_default_sort_order(Type::INT96, false),
+            SortOrder::INT96_TIMESTAMP
+        );
         assert_eq!(
             ColumnOrder::get_default_sort_order(Type::FLOAT, false),
             SortOrder::TOTAL_ORDER
@@ -2250,6 +2283,10 @@ mod tests {
             ColumnOrder::IEEE_754_TOTAL_ORDER.sort_order(),
             SortOrder::TOTAL_ORDER
         );
+        assert_eq!(
+            ColumnOrder::INT96_TIMESTAMP_ORDER.sort_order(),
+            SortOrder::INT96_TIMESTAMP
+        );
         assert_eq!(ColumnOrder::UNDEFINED.sort_order(), SortOrder::SIGNED);
     }
 

parquet/src/data_type.rs

@@ -143,12 +143,12 @@ impl PartialOrd for Int96 {
 impl Ord for Int96 {
     /// Order `Int96` correctly for (deprecated) timestamp types.
     ///
-    /// Note: this is done even though the Int96 type is deprecated and the
-    /// [spec does not define the sort order]
-    /// because some engines, notably Spark and Databricks Photon still write
-    /// Int96 timestamps and rely on their order for optimization.
+    /// Note: this is done even though the Int96 type is deprecated.
+    /// Because some engines, notably Spark and Databricks Photon, still write
+    /// Int96 timestamps, a new `ColumnOrder` variant has been added to
+    /// the Parquet specification. See [parquet-format/#584].
     ///
-    /// [spec does not define the sort order]: https://github.com/apache/parquet-format/blob/cf943c197f4fad826b14ba0c40eb0ffdab585285/src/main/thrift/parquet.thrift#L1079
+    /// [parquet-format/#584]: https://github.com/apache/parquet-format/pull/584
     fn cmp(&self, other: &Self) -> Ordering {
         match self.get_days().cmp(&other.get_days()) {
             Ordering::Equal => self.get_nanos().cmp(&other.get_nanos()),

parquet/src/file/writer.rs

@@ -1160,7 +1160,7 @@ mod tests {
     use crate::column::page::{Page, PageReader};
     use crate::column::reader::get_typed_column_reader;
     use crate::compression::{Codec, CodecOptionsBuilder, create_codec};
-    use crate::data_type::{BoolType, ByteArrayType, Int32Type};
+    use crate::data_type::{BoolType, ByteArrayType, Int32Type, Int96, Int96Type};
     use crate::file::page_index::column_index::ColumnIndexMetaData;
     use crate::file::properties::EnabledStatistics;
     use crate::file::serialized_reader::ReadOptionsBuilder;
@@ -2674,4 +2674,78 @@ mod tests {
         }
         writer.close().unwrap();
     }
+
+    #[test]
+    fn test_int96_interop() {
+        // this file has an INT96 column. rewrite it with min/max statistics sorted per
+        // recent changes to the spec. (see https://github.com/apache/parquet-format/pull/584)
+        let file = get_test_file("int96_timestamp_order.parquet");
+        let read_opts = ReadOptionsBuilder::new().with_page_index().build();
+        let reader = SerializedFileReader::new_with_options(file, read_opts).unwrap();
+        let file_metadata = reader.metadata().file_metadata();
+        let schema = file_metadata.schema_descr().root_schema_ptr();
+
+        // helper function to extract Int96 min/max from column metadata and the column index
+        fn retrieve_stats(metadata: &ParquetMetaData) -> (&[u8], &[u8], &Int96, &Int96) {
+            // sanity check that the proper column order is specified
+            let column_orders = metadata
+                .file_metadata()
+                .column_orders()
+                .expect("column_orders is missing");
+            assert_eq!(column_orders[0], ColumnOrder::INT96_TIMESTAMP_ORDER);
+
+            let stats = metadata
+                .row_group(0)
+                .column(0)
+                .statistics()
+                .expect("statistics missing");
+            let min = stats.min_bytes_opt().expect("min stats missing");
+            let max = stats.max_bytes_opt().expect("max stats missing");
+
+            let col_idx = metadata.column_index().expect("column index not present");
+            let col0 = match &col_idx[0][0] {
+                ColumnIndexMetaData::INT96(index) => index,
+                _ => panic!("expected INT96 stats"),
+            };
+            let col_min = col0.min_value(0).expect("ColumnIndex min not present");
+            let col_max = col0.max_value(0).expect("ColumnIndex max not present");
+
+            (min, max, col_min, col_max)
+        }
+
+        // save read stats for later
+        let (exp_min, exp_max, exp_col_min, exp_col_max) = retrieve_stats(reader.metadata());
+
+        // write file back out again
+        let props = Arc::new(WriterProperties::builder().build());
+        let output = Vec::<u8>::new();
+        let mut writer = SerializedFileWriter::new(output, schema, props).unwrap();
+
+        let mut rg_out = writer.next_row_group().unwrap();
+        let rg_in = reader.get_row_group(0).unwrap();
+
+        // int96 is column 0
+        let col_in = rg_in.get_column_reader(0).unwrap();
+        let mut typed_in = get_typed_column_reader::<Int96Type>(col_in);
+
+        let mut values = Vec::new();
+        typed_in.read_records(4, None, None, &mut values).unwrap();
+
+        let mut col_out = rg_out.next_column().unwrap().unwrap();
+        col_out
+            .typed::<Int96Type>()
+            .write_batch(&values, None, None)
+            .unwrap();
+        col_out.close().unwrap();
+        rg_out.close().unwrap();
+
+        let new_metadata = writer.close().unwrap();
+
+        // check that new stats match the original stats
+        let (new_min, new_max, new_col_min, new_col_max) = retrieve_stats(&new_metadata);
+        assert_eq!(new_min, exp_min);
+        assert_eq!(new_max, exp_max);
+        assert_eq!(new_col_min, exp_col_min);
+        assert_eq!(new_col_max, exp_col_max);
+    }
 }