builder

Author: cxzl25

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, TimestampPrecision};
+use crate::schema::{ArrowSchemaOptions, RootDataType, TimestampPrecision};
 use crate::stripe::{Stripe, StripeMetadata};
 
 const DEFAULT_BATCH_SIZE: usize = 8192;
@@ -149,10 +149,9 @@ impl<R> ArrowReaderBuilder<R> {
             .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_with_options(&metadata, self.timestamp_precision),
-            )
+            let options =
+                ArrowSchemaOptions::new().with_timestamp_precision(self.timestamp_precision);
+            Arc::new(projected_data_type.create_arrow_schema_with_options(&metadata, options))
         })
     }
 }

src/lib.rs

@@ -72,4 +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::TimestampPrecision;
+pub use schema::{ArrowSchemaOptions, TimestampPrecision};

src/schema.rs

@@ -37,6 +37,45 @@ pub enum TimestampPrecision {
     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.
@@ -73,22 +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, TimestampPrecision::default())
+        self.create_arrow_schema_with_options(user_metadata, ArrowSchemaOptions::new())
     }
 
-    /// Convert into an Arrow schema with specified timestamp precision.
+    /// Convert into an Arrow schema with custom options.
     pub fn create_arrow_schema_with_options(
         &self,
         user_metadata: &HashMap<String, String>,
-        timestamp_precision: TimestampPrecision,
+        options: ArrowSchemaOptions,
     ) -> Schema {
         let fields = self
             .children
             .iter()
             .map(|col| {
                 let dt = col
                     .data_type()
-                    .to_arrow_data_type_with_options(timestamp_precision);
+                    .to_arrow_data_type_with_options(options.clone());
                 Field::new(col.name(), dt, true)
             })
             .collect::<Vec<_>>();
@@ -455,14 +494,14 @@ 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(TimestampPrecision::default())
+        self.to_arrow_data_type_with_options(ArrowSchemaOptions::new())
     }
 
-    pub fn to_arrow_data_type_with_options(
-        &self,
-        timestamp_precision: TimestampPrecision,
-    ) -> ArrowDataType {
+    /// 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,
@@ -494,24 +533,24 @@ impl DataType {
                     .map(|col| {
                         let dt = col
                             .data_type()
-                            .to_arrow_data_type_with_options(timestamp_precision);
+                            .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_with_options(timestamp_precision);
+                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_with_options(timestamp_precision);
+                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_with_options(timestamp_precision);
+                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());
@@ -527,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_with_options(timestamp_precision);
+                        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));