parquet 0.5.9 → 0.5.10

data/ext/parquet/src/reader/common.rs

@@ -12,27 +12,81 @@ use magnus::value::ReprValue;
 use magnus::{Error as MagnusError, Ruby, Value};
 
 use crate::header_cache::StringCache;
+use crate::logger::RubyLogger;
 use crate::ruby_reader::{RubyReader, ThreadSafeRubyReader};
 use crate::types::{ParquetGemError, TryIntoValue};
 use crate::ColumnRecord;
 
-/// Opens a parquet file or IO-like object for reading
-///
-/// This function handles both file paths (as strings) and IO-like objects,
-/// returning either a File or a ThreadSafeRubyReader that can be used with
-/// parquet readers.
-pub fn open_parquet_source(
+use super::format_detector::{detect_file_format, detect_format_from_extension, FileFormat};
+
+/// Represents the different data sources we can open
+pub enum DataSource {
+    Parquet(Either<File, ThreadSafeRubyReader>),
+    Arrow(Either<File, ThreadSafeRubyReader>),
+}
+
+/// Opens a data file (Parquet or Arrow) for reading, automatically detecting the format
+pub fn open_data_source(
     ruby: Rc<Ruby>,
     to_read: Value,
-) -> Result<Either<File, ThreadSafeRubyReader>, ParquetGemError> {
+    ruby_logger: &RubyLogger,
+) -> Result<DataSource, ParquetGemError> {
     if to_read.is_kind_of(ruby.class_string()) {
         let path_string = to_read.to_r_string()?;
         let file_path = unsafe { path_string.as_str()? };
-        let file = File::open(file_path).map_err(ParquetGemError::from)?;
-        Ok(Either::Left(file))
+
+        // Try to detect format from extension first
+        let format_hint = detect_format_from_extension(file_path);
+
+        let mut file = File::open(file_path).map_err(ParquetGemError::from)?;
+
+        // Detect actual format from file content
+        let format = detect_file_format(&mut file)?;
+
+        // Warn if extension doesn't match content
+        if let Some(hint) = format_hint {
+            if hint != format {
+                ruby_logger.warn(|| {
+                    format!(
+                        "Extension implied format {:?} but actual format is {:?}",
+                        hint, format
+                    )
+                })?;
+            }
+        }
+
+        match format {
+            FileFormat::Parquet => Ok(DataSource::Parquet(Either::Left(file))),
+            FileFormat::Arrow => Ok(DataSource::Arrow(Either::Left(file))),
+        }
     } else {
-        let readable = ThreadSafeRubyReader::new(RubyReader::new(ruby, to_read)?);
-        Ok(Either::Right(readable))
+        // For IO-like objects, we need to use a temporary file
+        use std::io::{Read, Write};
+        use tempfile::NamedTempFile;
+
+        let mut readable = RubyReader::new(ruby.clone(), to_read)?;
+        let mut temp_file = NamedTempFile::new().map_err(ParquetGemError::from)?;
+
+        // Copy the entire content to the temporary file
+        let mut buffer = vec![0u8; 8192];
+        loop {
+            let bytes_read = readable.read(&mut buffer)?;
+            if bytes_read == 0 {
+                break;
+            }
+            temp_file.write_all(&buffer[..bytes_read])?;
+        }
+        temp_file.flush()?;
+
+        // Detect format from the temporary file
+        let mut file = temp_file.reopen()?;
+        let format = detect_file_format(&mut file)?;
+
+        // Use the temporary file as the source
+        match format {
+            FileFormat::Parquet => Ok(DataSource::Parquet(Either::Left(file))),
+            FileFormat::Arrow => Ok(DataSource::Arrow(Either::Left(file))),
+        }
     }
 }
 

data/ext/parquet/src/reader/format_detector.rs

@@ -0,0 +1,69 @@
+use crate::types::ParquetGemError;
+use std::io::{Read, Seek, SeekFrom};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum FileFormat {
+    Parquet,
+    Arrow,
+}
+
+/// Detect the file format by examining magic bytes
+pub fn detect_file_format<R: Read + Seek>(source: &mut R) -> Result<FileFormat, ParquetGemError> {
+    let mut magic = [0u8; 8];
+
+    // Read the first 8 bytes
+    let bytes_read = source.read(&mut magic).map_err(ParquetGemError::from)?;
+
+    // Reset to beginning
+    source
+        .seek(SeekFrom::Start(0))
+        .map_err(ParquetGemError::from)?;
+
+    if bytes_read >= 6 {
+        // Arrow IPC file format magic: "ARROW1\0\0"
+        if &magic[0..6] == b"ARROW1" {
+            return Ok(FileFormat::Arrow);
+        }
+    }
+
+    if bytes_read >= 4 {
+        // Parquet magic: "PAR1" at start
+        if &magic[0..4] == b"PAR1" {
+            return Ok(FileFormat::Parquet);
+        }
+    }
+
+    // If we can't detect from the beginning, check the end for Parquet
+    // Parquet files also have "PAR1" at the end
+    if let Ok(pos) = source.seek(SeekFrom::End(-4)) {
+        if pos >= 4 {
+            let mut end_magic = [0u8; 4];
+            if source.read_exact(&mut end_magic).is_ok() && &end_magic == b"PAR1" {
+                // Important: Reset to beginning before returning
+                source
+                    .seek(SeekFrom::Start(0))
+                    .map_err(ParquetGemError::from)?;
+                return Ok(FileFormat::Parquet);
+            }
+        }
+    }
+
+    // Always reset to beginning, even for unknown format
+    source
+        .seek(SeekFrom::Start(0))
+        .map_err(ParquetGemError::from)?;
+
+    Err(ParquetGemError::UnknownFormat)
+}
+
+/// Detect format from file extension as a fallback
+pub fn detect_format_from_extension(path: &str) -> Option<FileFormat> {
+    let lower = path.to_lowercase();
+    if lower.ends_with(".parquet") || lower.ends_with(".parq") {
+        Some(FileFormat::Parquet)
+    } else if lower.ends_with(".arrow") || lower.ends_with(".feather") || lower.ends_with(".ipc") {
+        Some(FileFormat::Arrow)
+    } else {
+        None
+    }
+}

data/ext/parquet/src/reader/mod.rs

@@ -1,4 +1,6 @@
+mod arrow_reader;
 mod common;
+mod format_detector;
 mod parquet_column_reader;
 mod parquet_row_reader;
 mod unified;
@@ -188,7 +190,10 @@ pub fn parse_metadata(_rb_self: Value, args: &[Value]) -> Result<Value, MagnusEr
     if args.len() != 1 {
         return Err(MagnusError::new(
             magnus::exception::arg_error(),
-            format!("metadata expects exactly 1 argument (file path or IO-like object), got {}", args.len()),
+            format!(
+                "metadata expects exactly 1 argument (file path or IO-like object), got {}",
+                args.len()
+            ),
         ));
     }
 
@@ -208,4 +213,4 @@ pub fn parse_metadata(_rb_self: Value, args: &[Value]) -> Result<Value, MagnusEr
     let metadata = reader.finish().map_err(ParquetGemError::Parquet)?;
 
     Ok(RubyParquetMetaData(metadata).try_into_value_with(&ruby)?)
-}
+}

data/ext/parquet/src/reader/unified/mod.rs

@@ -13,8 +13,11 @@ use std::collections::HashMap;
 use std::rc::Rc;
 use std::sync::OnceLock;
 
+use super::arrow_reader::{
+    process_arrow_column_data, process_arrow_file_column_data, process_arrow_row_data,
+};
 use super::common::{
-    create_batch_reader, handle_block_or_enum, handle_empty_file, open_parquet_source,
+    create_batch_reader, handle_block_or_enum, handle_empty_file, open_data_source, DataSource,
 };
 use crate::types::ArrayWrapper;
 
@@ -100,34 +103,99 @@ pub fn parse_parquet_unified(
         }
     }
 
-    // Open the Parquet source
-    let source = open_parquet_source(ruby.clone(), to_read)?;
+    // Open the data source and detect format
+    let source = open_data_source(ruby.clone(), to_read, &ruby_logger)?;
 
-    // Based on the parser type, handle the data differently
-    match parser_type {
-        ParserType::Row { strict } => {
-            // Handle row-based parsing
+    // Based on the source format and parser type, handle the data differently
+    match (source, &parser_type) {
+        (DataSource::Parquet(reader), ParserType::Row { strict }) => {
+            // Handle Parquet row-based parsing
             process_row_data(
                 ruby.clone(),
-                source,
+                reader,
                 &columns,
                 result_type,
-                strict,
+                *strict,
                 &ruby_logger,
             )?;
         }
-        ParserType::Column { batch_size, strict } => {
-            // Handle column-based parsing
+        (DataSource::Parquet(reader), ParserType::Column { batch_size, strict }) => {
+            // Handle Parquet column-based parsing
             process_column_data(
                 ruby.clone(),
-                source,
+                reader,
                 &columns,
                 result_type,
-                batch_size,
-                strict,
+                *batch_size,
+                *strict,
                 &ruby_logger,
             )?;
         }
+        (DataSource::Arrow(reader), ParserType::Row { strict }) => {
+            // Handle Arrow row-based parsing
+            match reader {
+                Either::Left(file) => {
+                    // For seekable files, use FileReader which handles IPC file format
+                    use arrow_ipc::reader::FileReader;
+                    let file_reader = FileReader::try_new(file, None)
+                        .map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                    use super::arrow_reader::process_arrow_file_row_data;
+                    process_arrow_file_row_data(
+                        ruby.clone(),
+                        file_reader,
+                        &columns,
+                        result_type,
+                        *strict,
+                        &ruby_logger,
+                    )?;
+                }
+                Either::Right(readable) => {
+                    use arrow_ipc::reader::StreamReader;
+                    let stream_reader = StreamReader::try_new(readable, None)
+                        .map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+                    process_arrow_row_data(
+                        ruby.clone(),
+                        stream_reader,
+                        &columns,
+                        result_type,
+                        *strict,
+                        &ruby_logger,
+                    )?;
+                }
+            }
+        }
+        (DataSource::Arrow(reader), ParserType::Column { batch_size, strict }) => {
+            // Handle Arrow column-based parsing
+            match reader {
+                Either::Left(file) => {
+                    // For seekable files, we can use the optimized FileReader
+                    process_arrow_file_column_data(
+                        ruby.clone(),
+                        file,
+                        &columns,
+                        result_type,
+                        *batch_size,
+                        *strict,
+                        &ruby_logger,
+                    )?;
+                }
+                Either::Right(readable) => {
+                    use arrow_ipc::reader::StreamReader;
+                    let stream_reader = StreamReader::try_new(readable, None)
+                        .map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+                    process_arrow_column_data(
+                        ruby.clone(),
+                        stream_reader,
+                        &columns,
+                        result_type,
+                        *batch_size,
+                        *strict,
+                        &ruby_logger,
+                    )?;
+                }
+            }
+        }
     }
 
     Ok(ruby.qnil().into_value_with(&ruby))

data/ext/parquet/src/types/core_types.rs

@@ -108,6 +108,7 @@ pub enum PrimitiveType {
     Float32,
     Float64,
     Decimal128(u8, i8),
+    Decimal256(u8, i8),
     Boolean,
     String,
     Binary,

data/ext/parquet/src/types/mod.rs

@@ -23,10 +23,11 @@ pub use writer_types::*;
 // Common imports used across the module
 use arrow_array::cast::downcast_array;
 use arrow_array::{
-    Array, BinaryArray, BooleanArray, Date32Array, Date64Array, Float16Array, Float32Array,
-    Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray, NullArray, StringArray,
-    StructArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
-    TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
+    Array, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal128Array, Decimal256Array,
+    Float16Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array,
+    ListArray, NullArray, StringArray, StructArray, TimestampMicrosecondArray,
+    TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt16Array,
+    UInt32Array, UInt64Array, UInt8Array,
 };
 use arrow_schema::{DataType, TimeUnit};
 use magnus::{value::ReprValue, Error as MagnusError, IntoValue, Ruby, Value};
@@ -54,6 +55,10 @@ pub enum ParquetGemError {
     Parquet(#[from] parquet::errors::ParquetError),
     #[error("Arrow error: {0}")]
     Arrow(#[from] arrow_schema::ArrowError),
+    #[error("Arrow IPC error: {0}")]
+    ArrowIpc(String),
+    #[error("Unknown file format")]
+    UnknownFormat,
     #[error("UTF-8 error: {0}")]
     Utf8Error(#[from] simdutf8::basic::Utf8Error),
     #[error("Jiff error: {0}")]
@@ -62,6 +67,8 @@ pub enum ParquetGemError {
     InvalidDecimal(String),
     #[error("Failed to parse UUID: {0}")]
     UuidError(#[from] uuid::Error),
+    #[error("Decimals larger than 128 bits are not supported")]
+    DecimalWouldBeTruncated,
 }
 
 #[derive(Debug)]