wip: date64 when coerce_types is false

dsgibbons · Aug 21, 2024 · a7f63ac · a7f63ac
1 parent 6c59b76
commit a7f63ac
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs
@@ -858,6 +858,13 @@ fn write_leaf(writer: &mut ColumnWriter<'_>, levels: &ArrayLevels) -> Result<usi
         }
         ColumnWriter::Int64ColumnWriter(ref mut typed) => {
             match column.data_type() {
+                ArrowDataType::Date64 => {
+                    let array = arrow_cast::cast(column, &ArrowDataType::Int64)?;
+
+                    let array = array.as_primitive::<Int64Type>();
+                    write_primitive(typed, array.values(), levels)
+                    // TODO: write test case to check - maybe use overflow condition 6 million years in future to check round trip
+                }
                 ArrowDataType::Int64 => {
                     let array = column.as_primitive::<Int64Type>();
                     write_primitive(typed, array.values(), levels)

diff --git a/parquet/src/arrow/schema/mod.rs b/parquet/src/arrow/schema/mod.rs
@@ -412,12 +412,17 @@ fn arrow_to_parquet_type(field: &Field) -> Result<Type> {
             .with_repetition(repetition)
             .with_id(id)
             .build(),
-        // date64 is cast to date32 (#1666)
-        DataType::Date64 => Type::primitive_type_builder(name, PhysicalType::INT32)
-            .with_logical_type(Some(LogicalType::Date))
-            .with_repetition(repetition)
-            .with_id(id)
-            .build(),
+        DataType::Date64 => {
+            let physical_type = match coerce_types {  // TODO: set via properties
+                true => PhysicalType::INT32,
+                false => PhysicalType::INT64,
+            };
+            Type::primitive_type_builder(name, physical_type)
+                .with_logical_type(Some(LogicalType::Date))
+                .with_repetition(repetition)
+                .with_id(id)
+                .build()
+        },
         DataType::Time32(TimeUnit::Second) => {
             // Cannot represent seconds in LogicalType
             Type::primitive_type_builder(name, PhysicalType::INT32)