fix timestamp 0001-01-01 overflow (#67)

cxzl25 · web-flow · commit 7ea0ce7de45f · 2025-11-17T11:47:15.000+08:00
* fix: timestamp 0001-01-01 overflow

* builder
diff --git a/src/array_decoder/timestamp.rs b/src/array_decoder/timestamp.rs
@@ -251,11 +251,28 @@ impl<T: ArrowTimestampType> ArrayBatchDecoder for TimestampOffsetArrayDecoder<T>
         let convert_timezone = |ts| {
             // Convert from writer timezone to reader timezone (which we default to UTC)
             // TODO: more efficient way of doing this?
-            self.writer_tz
-                .timestamp_nanos(ts)
-                .naive_local()
-                .and_utc()
-                .timestamp_nanos_opt()
+            let microseconds_in_timeunit = match T::UNIT {
+                TimeUnit::Second => 1_000_000,
+                TimeUnit::Millisecond => 1_000,
+                TimeUnit::Microsecond => 1,
+                TimeUnit::Nanosecond => -1, // not used in this case
+            };
+
+            match T::UNIT {
+                TimeUnit::Second | TimeUnit::Millisecond | TimeUnit::Microsecond => self
+                    .writer_tz
+                    .timestamp_micros(ts * microseconds_in_timeunit)
+                    .single()
+                    .map(|dt| {
+                        dt.naive_local().and_utc().timestamp_micros() / microseconds_in_timeunit
+                    }),
+                TimeUnit::Nanosecond => self
+                    .writer_tz
+                    .timestamp_nanos(ts)
+                    .naive_local()
+                    .and_utc()
+                    .timestamp_nanos_opt(),
+            }
         };
         let array = array
             // first try to convert all non-nullable batches to non-nullable batches
diff --git a/src/arrow_reader.rs b/src/arrow_reader.rs
@@ -29,7 +29,7 @@ use crate::projection::ProjectionMask;
 use crate::reader::metadata::{read_metadata, FileMetadata};
 use crate::reader::ChunkReader;
 use crate::row_selection::RowSelection;
-use crate::schema::RootDataType;
+use crate::schema::{ArrowSchemaOptions, RootDataType, TimestampPrecision};
 use crate::stripe::{Stripe, StripeMetadata};
 
 const DEFAULT_BATCH_SIZE: usize = 8192;
@@ -42,6 +42,7 @@ pub struct ArrowReaderBuilder<R> {
     pub(crate) schema_ref: Option<SchemaRef>,
     pub(crate) file_byte_range: Option<Range<usize>>,
     pub(crate) row_selection: Option<RowSelection>,
+    pub(crate) timestamp_precision: TimestampPrecision,
 }
 
 impl<R> ArrowReaderBuilder<R> {
@@ -54,6 +55,7 @@ impl<R> ArrowReaderBuilder<R> {
             schema_ref: None,
             file_byte_range: None,
             row_selection: None,
+            timestamp_precision: TimestampPrecision::default(),
         }
     }
 
@@ -109,6 +111,28 @@ impl<R> ArrowReaderBuilder<R> {
         self
     }
 
+    /// Sets the timestamp precision for reading timestamp columns.
+    ///
+    /// By default, timestamps are read as Nanosecond precision.
+    /// Use this method to switch to Microsecond precision if needed for compatibility.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use std::fs::File;
+    /// # use orc_rust::arrow_reader::ArrowReaderBuilder;
+    /// # use orc_rust::schema::TimestampPrecision;
+    /// let file = File::open("/path/to/file.orc").unwrap();
+    /// let reader = ArrowReaderBuilder::try_new(file)
+    ///     .unwrap()
+    ///     .with_timestamp_precision(TimestampPrecision::Microsecond)
+    ///     .build();
+    /// ```
+    pub fn with_timestamp_precision(mut self, precision: TimestampPrecision) -> Self {
+        self.timestamp_precision = precision;
+        self
+    }
+
     /// Returns the currently computed schema
     ///
     /// Unless [`with_schema`](Self::with_schema) was called, this is computed dynamically
@@ -124,9 +148,11 @@ impl<R> ArrowReaderBuilder<R> {
             .iter()
             .map(|(key, value)| (key.clone(), String::from_utf8_lossy(value).to_string()))
             .collect::<HashMap<_, _>>();
-        self.schema_ref
-            .clone()
-            .unwrap_or_else(|| Arc::new(projected_data_type.create_arrow_schema(&metadata)))
+        self.schema_ref.clone().unwrap_or_else(|| {
+            let options =
+                ArrowSchemaOptions::new().with_timestamp_precision(self.timestamp_precision);
+            Arc::new(projected_data_type.create_arrow_schema_with_options(&metadata, options))
+        })
     }
 }
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -72,3 +72,4 @@ pub use arrow_writer::{ArrowWriter, ArrowWriterBuilder};
 #[cfg(feature = "async")]
 pub use async_arrow_reader::ArrowStreamReader;
 pub use row_selection::{RowSelection, RowSelector};
+pub use schema::{ArrowSchemaOptions, TimestampPrecision};
diff --git a/src/schema.rs b/src/schema.rs
@@ -27,6 +27,55 @@ use crate::proto;
 
 use arrow::datatypes::{DataType as ArrowDataType, Field, Schema, TimeUnit, UnionMode};
 
+/// Configuration for timestamp precision when converting ORC timestamps to Arrow.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum TimestampPrecision {
+    /// Convert timestamps to microseconds (lower precision).
+    Microsecond,
+    /// Convert timestamps to nanoseconds (default, higher precision).
+    #[default]
+    Nanosecond,
+}
+
+/// Builder for configuring Arrow schema conversion options.
+#[derive(Debug, Clone)]
+pub struct ArrowSchemaOptions {
+    timestamp_precision: TimestampPrecision,
+}
+
+impl Default for ArrowSchemaOptions {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ArrowSchemaOptions {
+    /// Create a new options builder with default values.
+    /// - Timestamp precision is [`TimestampPrecision::Nanosecond`]
+    pub fn new() -> Self {
+        Self {
+            timestamp_precision: TimestampPrecision::default(),
+        }
+    }
+
+    /// Set the timestamp precision for converting ORC timestamps to Arrow.
+    ///
+    /// ORC timestamps have nanosecond precision, but you may want to convert
+    /// them to microseconds for compatibility with systems that don't support
+    /// nanosecond precision.
+    ///
+    /// Default: [`TimestampPrecision::Nanosecond`]
+    pub fn with_timestamp_precision(mut self, precision: TimestampPrecision) -> Self {
+        self.timestamp_precision = precision;
+        self
+    }
+
+    /// Get the timestamp precision
+    fn timestamp_precision(&self) -> TimestampPrecision {
+        self.timestamp_precision
+    }
+}
+
 /// Represents the root data type of the ORC file. Contains multiple named child types
 /// which map to the columns available. Allows projecting only specific columns from
 /// the base schema.
@@ -63,11 +112,22 @@ impl RootDataType {
 
     /// Convert into an Arrow schema.
     pub fn create_arrow_schema(&self, user_metadata: &HashMap<String, String>) -> Schema {
+        self.create_arrow_schema_with_options(user_metadata, ArrowSchemaOptions::new())
+    }
+
+    /// Convert into an Arrow schema with custom options.
+    pub fn create_arrow_schema_with_options(
+        &self,
+        user_metadata: &HashMap<String, String>,
+        options: ArrowSchemaOptions,
+    ) -> Schema {
         let fields = self
             .children
             .iter()
             .map(|col| {
-                let dt = col.data_type().to_arrow_data_type();
+                let dt = col
+                    .data_type()
+                    .to_arrow_data_type_with_options(options.clone());
                 Field::new(col.name(), dt, true)
             })
             .collect::<Vec<_>>();
@@ -434,7 +494,19 @@ impl DataType {
         Ok(dt)
     }
 
+    /// Convert this ORC data type to an Arrow data type with default options.
     pub fn to_arrow_data_type(&self) -> ArrowDataType {
+        self.to_arrow_data_type_with_options(ArrowSchemaOptions::new())
+    }
+
+    /// Convert this ORC data type to an Arrow data type with custom options.
+    pub fn to_arrow_data_type_with_options(&self, options: ArrowSchemaOptions) -> ArrowDataType {
+        let timestamp_precision = options.timestamp_precision();
+        let time_unit = match timestamp_precision {
+            TimestampPrecision::Microsecond => TimeUnit::Microsecond,
+            TimestampPrecision::Nanosecond => TimeUnit::Nanosecond,
+        };
+
         match self {
             DataType::Boolean { .. } => ArrowDataType::Boolean,
             DataType::Byte { .. } => ArrowDataType::Int8,
@@ -450,33 +522,35 @@ impl DataType {
             DataType::Decimal {
                 precision, scale, ..
             } => ArrowDataType::Decimal128(*precision as u8, *scale as i8), // TODO: safety of cast?
-            DataType::Timestamp { .. } => ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
+            DataType::Timestamp { .. } => ArrowDataType::Timestamp(time_unit, None),
             DataType::TimestampWithLocalTimezone { .. } => {
-                ArrowDataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into()))
+                ArrowDataType::Timestamp(time_unit, Some("UTC".into()))
             }
             DataType::Date { .. } => ArrowDataType::Date32,
             DataType::Struct { children, .. } => {
                 let children = children
                     .iter()
                     .map(|col| {
-                        let dt = col.data_type().to_arrow_data_type();
+                        let dt = col
+                            .data_type()
+                            .to_arrow_data_type_with_options(options.clone());
                         Field::new(col.name(), dt, true)
                     })
                     .collect();
                 ArrowDataType::Struct(children)
             }
             DataType::List { child, .. } => {
-                let child = child.to_arrow_data_type();
+                let child = child.to_arrow_data_type_with_options(options);
                 ArrowDataType::new_list(child, true)
             }
             DataType::Map { key, value, .. } => {
                 // TODO: this needs to be kept in sync with MapArrayDecoder
                 //       move to common location?
                 // TODO: should it be "keys" and "values" (like arrow-rs)
                 //       or "key" and "value" like PyArrow and in Schema.fbs?
-                let key = key.to_arrow_data_type();
+                let key = key.to_arrow_data_type_with_options(options.clone());
                 let key = Field::new("keys", key, false);
-                let value = value.to_arrow_data_type();
+                let value = value.to_arrow_data_type_with_options(options);
                 let value = Field::new("values", value, true);
 
                 let dt = ArrowDataType::Struct(vec![key, value].into());
@@ -492,7 +566,7 @@ impl DataType {
                         // TODO: Support up to including 256
                         //       Need to do Union within Union
                         let index = index as u8 as i8;
-                        let arrow_dt = variant.to_arrow_data_type();
+                        let arrow_dt = variant.to_arrow_data_type_with_options(options.clone());
                         // Name shouldn't matter here (only ORC struct types give names to subtypes anyway)
                         // Using naming convention following PyArrow for easier comparison
                         let field = Arc::new(Field::new(format!("_union_{index}"), arrow_dt, true));
diff --git a/tests/basic/main.rs b/tests/basic/main.rs
@@ -19,6 +19,7 @@ use std::fs::File;
 use std::ops::Range;
 use std::sync::Arc;
 
+use crate::misc::{LONG_BOOL_EXPECTED, LONG_STRING_DICT_EXPECTED, LONG_STRING_EXPECTED};
 use arrow::datatypes::{DataType, Decimal128Type, DecimalType, Field, Schema, TimeUnit};
 use arrow::record_batch::{RecordBatch, RecordBatchReader};
 use arrow::util::pretty;
@@ -28,8 +29,7 @@ use orc_rust::arrow_reader::{ArrowReader, ArrowReaderBuilder};
 #[cfg(feature = "async")]
 use orc_rust::async_arrow_reader::ArrowStreamReader;
 use orc_rust::projection::ProjectionMask;
-
-use crate::misc::{LONG_BOOL_EXPECTED, LONG_STRING_DICT_EXPECTED, LONG_STRING_EXPECTED};
+use orc_rust::TimestampPrecision;
 
 mod misc;
 
@@ -701,6 +701,55 @@ pub fn decimal128_timestamps_1900_test() {
     assert_batches_eq(&[batch], &expected);
 }
 
+#[test]
+pub fn timestamps_0001_test() {
+    let path = integration_path("timestamps_0001.orc");
+    let f = File::open(path).expect("no file found");
+    let mut reader = ArrowReaderBuilder::try_new(f)
+        .unwrap()
+        .with_schema(Arc::new(Schema::new(vec![Field::new(
+            "c1",
+            DataType::Timestamp(TimeUnit::Microsecond, None),
+            true,
+        )])))
+        .build();
+    let batch = reader.next().unwrap().unwrap();
+
+    let expected = [
+        "+---------------------+",
+        "| c1                  |",
+        "+---------------------+",
+        "| 0000-12-30T00:00:00 |",
+        "+---------------------+",
+    ];
+    assert_batches_eq(&[batch], &expected);
+}
+
+#[test]
+pub fn timestamps_0001_projection_test() {
+    let path = integration_path("timestamps_0001.orc");
+    let f = File::open(path).expect("no file found");
+
+    let builder = ArrowReaderBuilder::try_new(f).unwrap();
+    let projection = ProjectionMask::named_roots(builder.file_metadata().root_data_type(), &["c1"]);
+
+    let mut reader = builder
+        .with_projection(projection)
+        .with_timestamp_precision(TimestampPrecision::Microsecond)
+        .build();
+
+    let batch = reader.next().unwrap().unwrap();
+
+    let expected = [
+        "+---------------------+",
+        "| c1                  |",
+        "+---------------------+",
+        "| 0000-12-30T00:00:00 |",
+        "+---------------------+",
+    ];
+    assert_batches_eq(&[batch], &expected);
+}
+
 // From https://github.com/apache/arrow-rs/blob/7705acad845e8b2a366a08640f7acb4033ed7049/arrow-flight/src/sql/metadata/mod.rs#L67-L75
 pub fn assert_batches_eq(batches: &[RecordBatch], expected_lines: &[&str]) {
     let formatted = pretty::pretty_format_batches(batches).unwrap().to_string();
diff --git a/tests/integration/data/timestamps_0001.orc b/tests/integration/data/timestamps_0001.orc
