parquet 0.5.9 → 0.5.10

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f1a1f7c250b960dbe334145a537e865889fbc759e7b8bfbafcbadc77689972cd
-  data.tar.gz: 22116ec0b9fe89f0ad08a3674267bd00b141170b889091b476f3aab0d6be88a8
+  metadata.gz: 114891cfa5fa190e1f00d44803327f1c90cc11f64ba23f7f2a9cc9f9379da787
+  data.tar.gz: 9168b2be960faa93ce9c84d170c6e8f73819535bcedbf3d3b26869ff9829ecc6
 SHA512:
-  metadata.gz: ef8485d03247dd0d31993a774117669c1aaef5b875e7cb5c6f4c53e237a72fb81113ea35251426a21ea1ba24f8ae568bd2c3a158e6a45ce2416a308251d0f467
-  data.tar.gz: 672f38dfbf703dae996283fba8d137529e3089f569797df87feaac32fb0f956ea7c4d7ae57032715d1a21bd5bfa4dd728c2a3fe80174fb5fc0abdef51c73110a
+  metadata.gz: f07f99a188ac5fa0663616fba00b1990a2cbd6bb14462383915f0e1617c26c5ca481840c16179958f2b3760b334f176e2e4542d95e3cc922379948ac2b0bfa61
+  data.tar.gz: 42c7b0779d6e3fa46addc5fa92420f326418a54962d391e9b063db8378f8a5f8c2916b43f356649fc127e8fc582aa1e98d7afd71f0bc5f9700a0664ed46313f6

data/Cargo.lock

@@ -126,6 +126,7 @@ dependencies = [
  "arrow-data",
  "arrow-schema",
  "flatbuffers",
+ "lz4_flex",
 ]
 
 [[package]]
@@ -842,6 +843,8 @@ version = "0.1.0"
 dependencies = [
  "ahash",
  "arrow-array",
+ "arrow-buffer",
+ "arrow-ipc",
  "arrow-schema",
  "bytes",
  "either",

data/ext/parquet/Cargo.toml

@@ -12,6 +12,8 @@ rb-sys-env = "^0.2"
 [dependencies]
 ahash = "0.8"
 arrow-array = { git = "https://github.com/njaremko/arrow-rs", branch = "nathan/fix-time" }
+arrow-buffer = { git = "https://github.com/njaremko/arrow-rs", branch = "nathan/fix-time" }
+arrow-ipc = { git = "https://github.com/njaremko/arrow-rs", branch = "nathan/fix-time", features = ["lz4"] }
 arrow-schema = { git = "https://github.com/njaremko/arrow-rs", branch = "nathan/fix-time" }
 bytes = "^1.9"
 either = "1.9"

data/ext/parquet/build.rs

@@ -1,5 +1,5 @@
 pub fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let _rb_env = rb_sys_env::activate()?;
+    rb_sys_env::activate()?;
 
     Ok(())
 }

data/ext/parquet/src/lib.rs

@@ -19,6 +19,9 @@ use writer::write_rows;
 /// Initializes the Ruby extension and defines methods.
 #[magnus::init]
 fn init(ruby: &Ruby) -> Result<(), Error> {
+    // Require 'time' for Time.parse method
+    ruby.require("time")?;
+    
     let module = ruby.define_module("Parquet")?;
     module.define_module_function("metadata", magnus::method!(reader::parse_metadata, -1))?;
     module.define_module_function("each_row", magnus::method!(parse_parquet_rows, -1))?;

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

@@ -0,0 +1,579 @@
+use crate::header_cache::StringCache;
+use crate::logger::RubyLogger;
+use crate::types::ArrayWrapper;
+use crate::types::{
+    ColumnRecord, ParquetGemError, ParquetValueVec, ParserResultType, RowRecord, TryIntoValue,
+};
+use ahash::RandomState;
+use arrow_array::RecordBatch;
+use arrow_ipc::reader::{FileReader, StreamReader};
+use arrow_schema::Schema;
+use magnus::{Ruby, Value};
+use std::collections::HashMap;
+use std::fs::File;
+use std::io::Read;
+use std::rc::Rc;
+use std::sync::{Arc, OnceLock};
+
+/// Process Arrow IPC file data for column-based parsing
+pub fn process_arrow_column_data<R: Read>(
+    ruby: Rc<Ruby>,
+    reader: StreamReader<R>,
+    columns: &Option<Vec<String>>,
+    result_type: ParserResultType,
+    _batch_size: Option<usize>,
+    strict: bool,
+    ruby_logger: &RubyLogger,
+) -> Result<(), ParquetGemError> {
+    let schema = reader.schema();
+    ruby_logger.debug(|| format!("Arrow schema loaded: {:?}", schema))?;
+
+    // Filter schema if columns are specified
+    let _filtered_schema = if let Some(cols) = columns {
+        let mut fields = Vec::new();
+        for field in schema.fields() {
+            if cols.contains(&field.name().to_string()) {
+                fields.push(field.clone());
+            }
+        }
+        Arc::new(Schema::new(fields))
+    } else {
+        schema.clone()
+    };
+
+    match result_type {
+        ParserResultType::Hash => {
+            let headers = OnceLock::new();
+
+            for batch_result in reader {
+                let batch = batch_result.map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                // Filter columns if needed
+                let batch = if let Some(cols) = columns {
+                    filter_record_batch(&batch, cols)?
+                } else {
+                    batch
+                };
+
+                let local_headers = headers
+                    .get_or_init(|| {
+                        let schema = batch.schema();
+                        let fields = schema.fields();
+                        let mut header_string = Vec::with_capacity(fields.len());
+                        for field in fields {
+                            header_string.push(field.name().to_owned());
+                        }
+                        StringCache::intern_many(&header_string)
+                    })
+                    .as_ref()
+                    .map_err(|e| ParquetGemError::HeaderIntern(e.clone()))?;
+
+                let mut map =
+                    HashMap::with_capacity_and_hasher(local_headers.len(), RandomState::default());
+
+                batch
+                    .columns()
+                    .iter()
+                    .enumerate()
+                    .try_for_each(|(i, column)| {
+                        let header = local_headers[i];
+                        let values = ParquetValueVec::try_from(ArrayWrapper {
+                            array: column,
+                            strict,
+                        })?;
+                        map.insert(header, values.into_inner());
+                        Ok::<_, ParquetGemError>(())
+                    })?;
+
+                let record = ColumnRecord::Map::<RandomState>(map);
+                let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+            }
+        }
+        ParserResultType::Array => {
+            for batch_result in reader {
+                let batch = batch_result.map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                // Filter columns if needed
+                let batch = if let Some(cols) = columns {
+                    filter_record_batch(&batch, cols)?
+                } else {
+                    batch
+                };
+
+                let vec = batch
+                    .columns()
+                    .iter()
+                    .map(|column| {
+                        let values = ParquetValueVec::try_from(ArrayWrapper {
+                            array: column,
+                            strict,
+                        })?;
+                        Ok::<_, ParquetGemError>(values.into_inner())
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
+
+                let record = ColumnRecord::Vec::<RandomState>(vec);
+                let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
+/// Process Arrow IPC file data for row-based parsing
+pub fn process_arrow_row_data<R: Read>(
+    ruby: Rc<Ruby>,
+    reader: StreamReader<R>,
+    columns: &Option<Vec<String>>,
+    result_type: ParserResultType,
+    strict: bool,
+    ruby_logger: &RubyLogger,
+) -> Result<(), ParquetGemError> {
+    let schema = reader.schema();
+    ruby_logger.debug(|| format!("Arrow schema loaded: {:?}", schema))?;
+
+    match result_type {
+        ParserResultType::Hash => {
+            let headers = OnceLock::new();
+
+            for batch_result in reader {
+                let batch = batch_result.map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                // Filter columns if needed
+                let batch = if let Some(cols) = columns {
+                    filter_record_batch(&batch, cols)?
+                } else {
+                    batch
+                };
+
+                let local_headers = headers
+                    .get_or_init(|| {
+                        let schema = batch.schema();
+                        let fields = schema.fields();
+                        let mut header_string = Vec::with_capacity(fields.len());
+                        for field in fields {
+                            header_string.push(field.name().to_owned());
+                        }
+                        StringCache::intern_many(&header_string)
+                    })
+                    .as_ref()
+                    .map_err(|e| ParquetGemError::HeaderIntern(e.clone()))?;
+
+                // Convert columnar data to rows
+                for row_idx in 0..batch.num_rows() {
+                    let mut map = HashMap::with_capacity_and_hasher(
+                        local_headers.len(),
+                        RandomState::default(),
+                    );
+
+                    for (col_idx, column) in batch.columns().iter().enumerate() {
+                        let header = local_headers[col_idx];
+                        let value = extract_value_at_index(column, row_idx, strict)?;
+                        map.insert(header, value);
+                    }
+
+                    let record = RowRecord::Map::<RandomState>(map);
+                    let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+                }
+            }
+        }
+        ParserResultType::Array => {
+            for batch_result in reader {
+                let batch = batch_result.map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                // Filter columns if needed
+                let batch = if let Some(cols) = columns {
+                    filter_record_batch(&batch, cols)?
+                } else {
+                    batch
+                };
+
+                // Convert columnar data to rows
+                for row_idx in 0..batch.num_rows() {
+                    let mut row_vec = Vec::with_capacity(batch.num_columns());
+
+                    for column in batch.columns() {
+                        let value = extract_value_at_index(column, row_idx, strict)?;
+                        row_vec.push(value);
+                    }
+
+                    let record = RowRecord::Vec::<RandomState>(row_vec);
+                    let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+                }
+            }
+        }
+    }
+
+    Ok(())
+}
+
+/// Process Arrow IPC file with FileReader for row-based parsing
+pub fn process_arrow_file_row_data(
+    ruby: Rc<Ruby>,
+    reader: FileReader<File>,
+    columns: &Option<Vec<String>>,
+    result_type: ParserResultType,
+    strict: bool,
+    ruby_logger: &RubyLogger,
+) -> Result<(), ParquetGemError> {
+    let schema = reader.schema();
+    ruby_logger.debug(|| format!("Arrow file schema loaded: {:?}", schema))?;
+
+    match result_type {
+        ParserResultType::Hash => {
+            let headers = OnceLock::new();
+
+            for batch_result in reader {
+                let batch = batch_result.map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                // Filter columns if needed
+                let batch = if let Some(cols) = columns {
+                    filter_record_batch(&batch, cols)?
+                } else {
+                    batch
+                };
+
+                let local_headers = headers
+                    .get_or_init(|| {
+                        let schema = batch.schema();
+                        let fields = schema.fields();
+                        let mut header_string = Vec::with_capacity(fields.len());
+                        for field in fields {
+                            header_string.push(field.name().to_owned());
+                        }
+                        StringCache::intern_many(&header_string)
+                    })
+                    .as_ref()
+                    .map_err(|e| ParquetGemError::HeaderIntern(e.clone()))?;
+
+                // Convert columnar data to rows
+                for row_idx in 0..batch.num_rows() {
+                    let mut map = HashMap::with_capacity_and_hasher(
+                        local_headers.len(),
+                        RandomState::default(),
+                    );
+
+                    for (col_idx, column) in batch.columns().iter().enumerate() {
+                        let header = local_headers[col_idx];
+                        let value = extract_value_at_index(column, row_idx, strict)?;
+                        map.insert(header, value);
+                    }
+
+                    let record = RowRecord::Map::<RandomState>(map);
+                    let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+                }
+            }
+        }
+        ParserResultType::Array => {
+            for batch_result in reader {
+                let batch = batch_result.map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                // Filter columns if needed
+                let batch = if let Some(cols) = columns {
+                    filter_record_batch(&batch, cols)?
+                } else {
+                    batch
+                };
+
+                // Convert columnar data to rows
+                for row_idx in 0..batch.num_rows() {
+                    let mut row_vec = Vec::with_capacity(batch.num_columns());
+
+                    for column in batch.columns() {
+                        let value = extract_value_at_index(column, row_idx, strict)?;
+                        row_vec.push(value);
+                    }
+
+                    let record = RowRecord::Vec::<RandomState>(row_vec);
+                    let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+                }
+            }
+        }
+    }
+
+    Ok(())
+}
+
+/// Process Arrow IPC file with FileReader (for seekable sources)
+pub fn process_arrow_file_column_data(
+    ruby: Rc<Ruby>,
+    file: File,
+    columns: &Option<Vec<String>>,
+    result_type: ParserResultType,
+    _batch_size: Option<usize>,
+    strict: bool,
+    ruby_logger: &RubyLogger,
+) -> Result<(), ParquetGemError> {
+    let reader =
+        FileReader::try_new(file, None).map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+    let schema = reader.schema();
+    ruby_logger.debug(|| format!("Arrow file schema loaded: {:?}", schema))?;
+
+    // FileReader implements Iterator<Item = Result<RecordBatch, ArrowError>>
+    match result_type {
+        ParserResultType::Hash => {
+            let headers = OnceLock::new();
+
+            for batch_result in reader {
+                let batch = batch_result.map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                // Filter columns if needed
+                let batch = if let Some(cols) = columns {
+                    filter_record_batch(&batch, cols)?
+                } else {
+                    batch
+                };
+
+                let local_headers = headers
+                    .get_or_init(|| {
+                        let schema = batch.schema();
+                        let fields = schema.fields();
+                        let mut header_string = Vec::with_capacity(fields.len());
+                        for field in fields {
+                            header_string.push(field.name().to_owned());
+                        }
+                        StringCache::intern_many(&header_string)
+                    })
+                    .as_ref()
+                    .map_err(|e| ParquetGemError::HeaderIntern(e.clone()))?;
+
+                let mut map =
+                    HashMap::with_capacity_and_hasher(local_headers.len(), RandomState::default());
+
+                batch
+                    .columns()
+                    .iter()
+                    .enumerate()
+                    .try_for_each(|(i, column)| {
+                        let header = local_headers[i];
+                        let values = ParquetValueVec::try_from(ArrayWrapper {
+                            array: column,
+                            strict,
+                        })?;
+                        map.insert(header, values.into_inner());
+                        Ok::<_, ParquetGemError>(())
+                    })?;
+
+                let record = ColumnRecord::Map::<RandomState>(map);
+                let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+            }
+        }
+        ParserResultType::Array => {
+            for batch_result in reader {
+                let batch = batch_result.map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))?;
+
+                // Filter columns if needed
+                let batch = if let Some(cols) = columns {
+                    filter_record_batch(&batch, cols)?
+                } else {
+                    batch
+                };
+
+                let vec = batch
+                    .columns()
+                    .iter()
+                    .map(|column| {
+                        let values = ParquetValueVec::try_from(ArrayWrapper {
+                            array: column,
+                            strict,
+                        })?;
+                        Ok::<_, ParquetGemError>(values.into_inner())
+                    })
+                    .collect::<Result<Vec<_>, _>>()?;
+
+                let record = ColumnRecord::Vec::<RandomState>(vec);
+                let _: Value = ruby.yield_value(record.try_into_value_with(&ruby)?)?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
+/// Extract a single value from an Arrow array at a specific index
+fn extract_value_at_index(
+    array: &Arc<dyn arrow_array::Array>,
+    index: usize,
+    strict: bool,
+) -> Result<crate::types::ParquetField, ParquetGemError> {
+    use crate::types::ParquetField;
+    use arrow_array::*;
+    use arrow_schema::DataType;
+    use parquet::record::Field;
+
+    // Convert Arrow array value at index to Parquet Field
+    let field = match array.data_type() {
+        DataType::Boolean => {
+            let arr = array.as_any().downcast_ref::<BooleanArray>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Bool(arr.value(index))
+            }
+        }
+        DataType::Int8 => {
+            let arr = array.as_any().downcast_ref::<Int8Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Byte(arr.value(index) as i8)
+            }
+        }
+        DataType::Int16 => {
+            let arr = array.as_any().downcast_ref::<Int16Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Short(arr.value(index))
+            }
+        }
+        DataType::Int32 => {
+            let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Int(arr.value(index))
+            }
+        }
+        DataType::Int64 => {
+            let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Long(arr.value(index))
+            }
+        }
+        DataType::UInt8 => {
+            let arr = array.as_any().downcast_ref::<UInt8Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::UByte(arr.value(index))
+            }
+        }
+        DataType::UInt16 => {
+            let arr = array.as_any().downcast_ref::<UInt16Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::UShort(arr.value(index))
+            }
+        }
+        DataType::UInt32 => {
+            let arr = array.as_any().downcast_ref::<UInt32Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::UInt(arr.value(index))
+            }
+        }
+        DataType::UInt64 => {
+            let arr = array.as_any().downcast_ref::<UInt64Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::ULong(arr.value(index))
+            }
+        }
+        DataType::Float32 => {
+            let arr = array.as_any().downcast_ref::<Float32Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Float(arr.value(index))
+            }
+        }
+        DataType::Float64 => {
+            let arr = array.as_any().downcast_ref::<Float64Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Double(arr.value(index))
+            }
+        }
+        DataType::Utf8 => {
+            let arr = array.as_any().downcast_ref::<StringArray>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Str(arr.value(index).to_string())
+            }
+        }
+        DataType::Binary => {
+            let arr = array.as_any().downcast_ref::<BinaryArray>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Bytes(arr.value(index).into())
+            }
+        }
+        DataType::Date32 => {
+            let arr = array.as_any().downcast_ref::<Date32Array>().unwrap();
+            if arr.is_null(index) {
+                Field::Null
+            } else {
+                Field::Date(arr.value(index))
+            }
+        }
+        DataType::Timestamp(unit, _tz) => match unit {
+            arrow_schema::TimeUnit::Millisecond => {
+                let arr = array
+                    .as_any()
+                    .downcast_ref::<TimestampMillisecondArray>()
+                    .unwrap();
+                if arr.is_null(index) {
+                    Field::Null
+                } else {
+                    Field::TimestampMillis(arr.value(index))
+                }
+            }
+            arrow_schema::TimeUnit::Microsecond => {
+                let arr = array
+                    .as_any()
+                    .downcast_ref::<TimestampMicrosecondArray>()
+                    .unwrap();
+                if arr.is_null(index) {
+                    Field::Null
+                } else {
+                    Field::TimestampMicros(arr.value(index))
+                }
+            }
+            _ => Field::Null,
+        },
+        // Add more type handling as needed
+        _ => Field::Null,
+    };
+
+    // For Arrow files, we don't have Parquet logical types, so we use defaults
+    Ok(ParquetField {
+        field,
+        converted_type: parquet::basic::ConvertedType::NONE,
+        logical_type: None,
+        strict,
+    })
+}
+
+/// Filter a RecordBatch to only include specified columns
+fn filter_record_batch(
+    batch: &RecordBatch,
+    columns: &[String],
+) -> Result<RecordBatch, ParquetGemError> {
+    let schema = batch.schema();
+    let mut indices = Vec::new();
+    let mut fields = Vec::new();
+
+    for (i, field) in schema.fields().iter().enumerate() {
+        if columns.contains(&field.name().to_string()) {
+            indices.push(i);
+            fields.push(field.clone());
+        }
+    }
+
+    let new_schema = Arc::new(Schema::new(fields));
+    let new_columns: Vec<_> = indices.iter().map(|&i| batch.column(i).clone()).collect();
+
+    RecordBatch::try_new(new_schema, new_columns)
+        .map_err(|e| ParquetGemError::ArrowIpc(e.to_string()))
+}