parquet 0.2.12-arm64-darwin

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

@@ -0,0 +1,214 @@
+use crate::header_cache::StringCache;
+use crate::{
+    create_column_enumerator, utils::*, ColumnEnumeratorArgs, ColumnRecord, ForgottenFileHandle,
+    ParquetValueVec, ParserResultType, SeekableRubyValue,
+};
+use ahash::RandomState;
+use magnus::rb_sys::AsRawValue;
+use magnus::value::{Opaque, ReprValue};
+use magnus::IntoValue;
+use magnus::{Error as MagnusError, Ruby, Value};
+use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+use parquet::arrow::ProjectionMask;
+use std::collections::HashMap;
+use std::fs::File;
+use std::mem::ManuallyDrop;
+use std::os::fd::FromRawFd;
+use std::sync::OnceLock;
+
+use super::ReaderError;
+
+#[inline]
+pub fn parse_parquet_columns<'a>(rb_self: Value, args: &[Value]) -> Result<Value, MagnusError> {
+    let ruby = unsafe { Ruby::get_unchecked() };
+
+    let ParquetColumnsArgs {
+        to_read,
+        result_type,
+        columns,
+        batch_size,
+    } = parse_parquet_columns_args(&ruby, args)?;
+
+    if !ruby.block_given() {
+        return create_column_enumerator(ColumnEnumeratorArgs {
+            rb_self,
+            to_read,
+            result_type,
+            columns,
+            batch_size,
+        })
+        .map(|yield_enum| yield_enum.into_value_with(&ruby));
+    }
+
+    let (batch_reader, schema, num_rows) = 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(|e| ReaderError::FileOpen(e))?;
+
+        let mut builder =
+            ParquetRecordBatchReaderBuilder::try_new(file).map_err(|e| ReaderError::Parquet(e))?;
+        let schema = builder.schema().clone();
+        let num_rows = builder.metadata().file_metadata().num_rows();
+
+        // If columns are specified, project only those columns
+        if let Some(cols) = &columns {
+            // Get the parquet schema
+            let parquet_schema = builder.parquet_schema();
+
+            // Create a projection mask from column names
+            let projection =
+                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
+
+            builder = builder.with_projection(projection);
+        }
+
+        if let Some(batch_size) = batch_size {
+            builder = builder.with_batch_size(batch_size);
+        }
+
+        let reader = builder.build().map_err(|e| ReaderError::Parquet(e))?;
+
+        (reader, schema, num_rows)
+    } else if to_read.is_kind_of(ruby.class_io()) {
+        let raw_value = to_read.as_raw();
+        let fd = std::panic::catch_unwind(|| unsafe { rb_sys::rb_io_descriptor(raw_value) })
+            .map_err(|_| {
+                ReaderError::FileDescriptor("Failed to get file descriptor".to_string())
+            })?;
+
+        if fd < 0 {
+            return Err(ReaderError::InvalidFileDescriptor.into());
+        }
+
+        let file = unsafe { File::from_raw_fd(fd) };
+        let file = ForgottenFileHandle(ManuallyDrop::new(file));
+
+        let mut builder =
+            ParquetRecordBatchReaderBuilder::try_new(file).map_err(|e| ReaderError::Parquet(e))?;
+        let schema = builder.schema().clone();
+        let num_rows = builder.metadata().file_metadata().num_rows();
+
+        // If columns are specified, project only those columns
+        if let Some(cols) = &columns {
+            // Get the parquet schema
+            let parquet_schema = builder.parquet_schema();
+
+            // Create a projection mask from column names
+            let projection =
+                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
+
+            builder = builder.with_projection(projection);
+        }
+
+        if let Some(batch_size) = batch_size {
+            builder = builder.with_batch_size(batch_size);
+        }
+
+        let reader = builder.build().map_err(|e| ReaderError::Parquet(e))?;
+
+        (reader, schema, num_rows)
+    } else {
+        let readable = SeekableRubyValue(Opaque::from(to_read));
+
+        let mut builder = ParquetRecordBatchReaderBuilder::try_new(readable)
+            .map_err(|e| ReaderError::Parquet(e))?;
+        let schema = builder.schema().clone();
+        let num_rows = builder.metadata().file_metadata().num_rows();
+
+        // If columns are specified, project only those columns
+        if let Some(cols) = &columns {
+            // Get the parquet schema
+            let parquet_schema = builder.parquet_schema();
+
+            // Create a projection mask from column names
+            let projection =
+                ProjectionMask::columns(parquet_schema, cols.iter().map(|s| s.as_str()));
+
+            builder = builder.with_projection(projection);
+        }
+
+        if let Some(batch_size) = batch_size {
+            builder = builder.with_batch_size(batch_size);
+        }
+
+        let reader = builder.build().map_err(|e| ReaderError::Parquet(e))?;
+
+        (reader, schema, num_rows)
+    };
+
+    if num_rows == 0 {
+        let mut map =
+            HashMap::with_capacity_and_hasher(schema.fields().len(), RandomState::default());
+        let headers: Vec<String> = schema
+            .fields()
+            .iter()
+            .map(|field| field.name().to_string())
+            .collect();
+        let interned_headers =
+            StringCache::intern_many(&headers).map_err(|e| ReaderError::HeaderIntern(e))?;
+        for field in interned_headers.iter() {
+            map.insert(*field, vec![]);
+        }
+        let record = ColumnRecord::Map(map);
+        let _: Value = ruby.yield_value(record)?;
+        return Ok(ruby.qnil().into_value_with(&ruby));
+    }
+
+    match result_type {
+        ParserResultType::Hash => {
+            let headers = OnceLock::new();
+            let headers_clone = headers.clone();
+            let iter = batch_reader.map(move |batch| {
+                batch.map_err(ReaderError::Arrow).and_then(|batch| {
+                    let headers = headers_clone.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).unwrap()
+                    });
+
+                    let mut map =
+                        HashMap::with_capacity_and_hasher(headers.len(), RandomState::default());
+
+                    batch.columns().iter().enumerate().for_each(|(i, column)| {
+                        let header = headers[i];
+                        let values = ParquetValueVec::try_from(column.clone()).unwrap();
+                        map.insert(header, values.into_inner());
+                    });
+
+                    Ok(ColumnRecord::Map::<RandomState>(map))
+                })
+            });
+
+            for result in iter {
+                let record = result?;
+                let _: Value = ruby.yield_value(record)?;
+            }
+        }
+        ParserResultType::Array => {
+            let iter = batch_reader.map(|batch| {
+                batch.map_err(ReaderError::Arrow).and_then(|batch| {
+                    let vec = batch
+                        .columns()
+                        .into_iter()
+                        .map(|column| {
+                            let values = ParquetValueVec::try_from(column.clone()).unwrap();
+                            values.into_inner()
+                        })
+                        .collect();
+                    Ok(ColumnRecord::Vec::<RandomState>(vec))
+                })
+            });
+
+            for result in iter {
+                let record = result?;
+                let _: Value = ruby.yield_value(record)?;
+            }
+        }
+    }
+
+    Ok(ruby.qnil().into_value_with(&ruby))
+}

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

@@ -0,0 +1,157 @@
+use crate::header_cache::StringCache;
+use crate::{
+    create_row_enumerator, utils::*, ForgottenFileHandle, ParquetField, ParserResultType,
+    ReaderError, RowEnumeratorArgs, RowRecord, SeekableRubyValue,
+};
+use ahash::RandomState;
+use magnus::rb_sys::AsRawValue;
+use magnus::value::{Opaque, ReprValue};
+use magnus::IntoValue;
+use magnus::{Error as MagnusError, Ruby, Value};
+use parquet::file::reader::{FileReader, SerializedFileReader};
+use parquet::record::reader::RowIter as ParquetRowIter;
+use parquet::schema::types::{Type as SchemaType, TypePtr};
+use std::collections::HashMap;
+use std::fs::File;
+use std::mem::ManuallyDrop;
+use std::os::fd::FromRawFd;
+use std::sync::OnceLock;
+
+#[inline]
+pub fn parse_parquet_rows<'a>(rb_self: Value, args: &[Value]) -> Result<Value, MagnusError> {
+    let ruby = unsafe { Ruby::get_unchecked() };
+
+    let ParquetRowsArgs {
+        to_read,
+        result_type,
+        columns,
+    } = parse_parquet_rows_args(&ruby, args)?;
+
+    if !ruby.block_given() {
+        return create_row_enumerator(RowEnumeratorArgs {
+            rb_self,
+            to_read,
+            result_type,
+            columns,
+        })
+        .map(|yield_enum| yield_enum.into_value_with(&ruby));
+    }
+
+    let (schema, mut iter) = 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).unwrap();
+        let reader = SerializedFileReader::new(file).unwrap();
+        let schema = reader.metadata().file_metadata().schema().clone();
+
+        (schema, ParquetRowIter::from_file_into(Box::new(reader)))
+    } else if to_read.is_kind_of(ruby.class_io()) {
+        let raw_value = to_read.as_raw();
+        let fd = std::panic::catch_unwind(|| unsafe { rb_sys::rb_io_descriptor(raw_value) })
+            .map_err(|_| {
+                ReaderError::FileDescriptor("Failed to get file descriptor".to_string())
+            })?;
+
+        if fd < 0 {
+            return Err(ReaderError::InvalidFileDescriptor.into());
+        }
+
+        let file = unsafe { File::from_raw_fd(fd) };
+        let file = ForgottenFileHandle(ManuallyDrop::new(file));
+        let reader = SerializedFileReader::new(file).unwrap();
+        let schema = reader.metadata().file_metadata().schema().clone();
+
+        (schema, ParquetRowIter::from_file_into(Box::new(reader)))
+    } else {
+        let readable = SeekableRubyValue(Opaque::from(to_read));
+        let reader = SerializedFileReader::new(readable).unwrap();
+        let schema = reader.metadata().file_metadata().schema().clone();
+
+        (schema, ParquetRowIter::from_file_into(Box::new(reader)))
+    };
+
+    if let Some(cols) = columns {
+        let projection = create_projection_schema(&schema, &cols);
+        iter = iter.project(Some(projection.to_owned())).map_err(|e| {
+            MagnusError::new(
+                ruby.exception_runtime_error(),
+                format!("Failed to create projection: {}", e),
+            )
+        })?;
+    }
+
+    match result_type {
+        ParserResultType::Hash => {
+            let headers = OnceLock::new();
+            let headers_clone = headers.clone();
+            let iter = iter.map(move |row| {
+                row.and_then(|row| {
+                    let headers = headers_clone.get_or_init(|| {
+                        let column_count = row.get_column_iter().count();
+
+                        let mut header_string = Vec::with_capacity(column_count);
+                        for (k, _) in row.get_column_iter() {
+                            header_string.push(k.to_owned());
+                        }
+
+                        let headers = StringCache::intern_many(&header_string).unwrap();
+
+                        headers
+                    });
+
+                    let mut map =
+                        HashMap::with_capacity_and_hasher(headers.len(), RandomState::default());
+                    row.get_column_iter().enumerate().for_each(|(i, (_, v))| {
+                        map.insert(headers[i], ParquetField(v.clone()));
+                    });
+                    Ok(map)
+                })
+                .and_then(|row| Ok(RowRecord::Map::<RandomState>(row)))
+                .map_err(|e| ReaderError::Parquet(e))
+            });
+
+            for result in iter {
+                let record = result?;
+                let _: Value = ruby.yield_value(record)?;
+            }
+        }
+        ParserResultType::Array => {
+            let iter = iter.map(|row| {
+                row.and_then(|row| {
+                    let column_count = row.get_column_iter().count();
+                    let mut vec = Vec::with_capacity(column_count);
+                    row.get_column_iter()
+                        .for_each(|(_, v)| vec.push(ParquetField(v.clone())));
+                    Ok(vec)
+                })
+                .and_then(|row| Ok(RowRecord::Vec::<RandomState>(row)))
+                .map_err(|e| ReaderError::Parquet(e))
+            });
+
+            for result in iter {
+                let record = result?;
+                let _: Value = ruby.yield_value(record)?;
+            }
+        }
+    }
+
+    Ok(ruby.qnil().into_value_with(&ruby))
+}
+
+fn create_projection_schema(schema: &SchemaType, columns: &[String]) -> SchemaType {
+    if let SchemaType::GroupType { fields, .. } = schema {
+        let projected_fields: Vec<TypePtr> = fields
+            .iter()
+            .filter(|field| columns.contains(&field.name().to_string()))
+            .cloned()
+            .collect();
+
+        SchemaType::GroupType {
+            basic_info: schema.get_basic_info().clone(),
+            fields: projected_fields,
+        }
+    } else {
+        // Return original schema if not a group type
+        schema.clone()
+    }
+}

data/ext/parquet/src/ruby_integration.rs

@@ -0,0 +1,77 @@
+use std::{
+    fs::File,
+    io::{BufReader, SeekFrom},
+    mem::ManuallyDrop,
+};
+
+use bytes::Bytes;
+use magnus::{value::Opaque, Ruby, Value};
+use parquet::{
+    errors::ParquetError,
+    file::reader::{ChunkReader, Length},
+};
+use std::io::Read;
+
+use crate::ruby_reader::{build_ruby_reader, SeekableRead};
+
+const READ_BUFFER_SIZE: usize = 16 * 1024;
+
+pub struct SeekableRubyValue(pub Opaque<Value>);
+
+impl Length for SeekableRubyValue {
+    fn len(&self) -> u64 {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        let current_pos = reader.seek(SeekFrom::Current(0)).unwrap();
+        let file_len = reader.seek(SeekFrom::End(0)).unwrap();
+        reader.seek(SeekFrom::Start(current_pos)).unwrap();
+        file_len
+    }
+}
+
+impl ChunkReader for SeekableRubyValue {
+    type T = BufReader<Box<dyn SeekableRead>>;
+
+    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        reader.seek(SeekFrom::Start(start))?;
+        Ok(BufReader::with_capacity(READ_BUFFER_SIZE, reader))
+    }
+
+    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
+        let ruby = unsafe { Ruby::get_unchecked() };
+        let mut buffer = Vec::with_capacity(length);
+        let mut reader = build_ruby_reader(&ruby, ruby.get_inner(self.0)).unwrap();
+        reader.seek(SeekFrom::Start(start))?;
+        let read = reader.take(length as _).read_to_end(&mut buffer)?;
+
+        if read != length {
+            return Err(ParquetError::EOF(format!(
+                "Expected to read {} bytes, read only {}",
+                length, read
+            )));
+        }
+        Ok(buffer.into())
+    }
+}
+
+pub struct ForgottenFileHandle(pub ManuallyDrop<File>);
+
+impl Length for ForgottenFileHandle {
+    fn len(&self) -> u64 {
+        self.0.len()
+    }
+}
+
+impl ChunkReader for ForgottenFileHandle {
+    type T = BufReader<File>;
+
+    fn get_read(&self, start: u64) -> parquet::errors::Result<Self::T> {
+        self.0.get_read(start)
+    }
+
+    fn get_bytes(&self, start: u64, length: usize) -> parquet::errors::Result<Bytes> {
+        self.0.get_bytes(start, length)
+    }
+}

data/ext/parquet/src/ruby_reader.rs

@@ -0,0 +1,171 @@
+use magnus::{
+    value::{Opaque, ReprValue},
+    RClass, RString, Ruby, Value,
+};
+use std::io::{self, Read, Seek, SeekFrom, Write};
+use std::sync::OnceLock;
+
+static STRING_IO_CLASS: OnceLock<Opaque<RClass>> = OnceLock::new();
+
+/// A reader that can handle various Ruby input types (String, StringIO, IO-like objects)
+/// and provide a standard Read implementation for them.
+pub struct RubyReader<T> {
+    inner: T,
+    offset: usize,
+}
+
+pub trait SeekableRead: Read + Seek {}
+impl<T: Read + Seek> SeekableRead for T {}
+
+pub fn build_ruby_reader(
+    ruby: &Ruby,
+    input: Value,
+) -> Result<Box<dyn SeekableRead>, magnus::Error> {
+    if RubyReader::is_string_io(ruby, &input) {
+        RubyReader::from_string_io(ruby, input)
+    } else if RubyReader::is_io_like(&input) {
+        RubyReader::from_io(input)
+    } else {
+        RubyReader::from_string_like(input)
+    }
+}
+
+impl Seek for RubyReader<Value> {
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        let (whence, offset) = match pos {
+            SeekFrom::Start(i) => (0, i as i64),
+            SeekFrom::Current(i) => (1, i),
+            SeekFrom::End(i) => (2, i),
+        };
+
+        let new_position = self
+            .inner
+            .funcall("seek", (offset, whence))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        Ok(new_position)
+    }
+}
+
+impl Write for RubyReader<Value> {
+    fn write(&mut self, buf: &[u8]) -> Result<usize, io::Error> {
+        let ruby_bytes = RString::from_slice(buf);
+
+        let bytes_written = self
+            .inner
+            .funcall::<_, _, usize>("write", (ruby_bytes,))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        Ok(bytes_written)
+    }
+
+    fn flush(&mut self) -> Result<(), io::Error> {
+        self.inner
+            .funcall::<_, _, Value>("flush", ())
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        Ok(())
+    }
+}
+
+impl Seek for RubyReader<RString> {
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        match pos {
+            io::SeekFrom::Start(offset) => self.offset = offset as usize,
+            io::SeekFrom::Current(offset) => self.offset = (self.offset as i64 + offset) as usize,
+            io::SeekFrom::End(offset) => self.offset = self.inner.len() - offset as usize,
+        }
+        Ok(self.offset as u64)
+    }
+}
+
+impl RubyReader<Value> {
+    fn from_io(input: Value) -> Result<Box<dyn SeekableRead>, magnus::Error> {
+        if Self::is_io_like(&input) {
+            Ok(Box::new(Self::from_io_like(input)))
+        } else {
+            Err(magnus::Error::new(
+                magnus::exception::type_error(),
+                "Input is not an IO-like object",
+            ))
+        }
+    }
+
+    fn is_io_like(input: &Value) -> bool {
+        input.respond_to("read", false).unwrap_or(false)
+    }
+
+    fn from_io_like(input: Value) -> Self {
+        Self {
+            inner: input,
+            offset: 0,
+        }
+    }
+}
+
+impl RubyReader<RString> {
+    pub fn from_string_io(
+        ruby: &Ruby,
+        input: Value,
+    ) -> Result<Box<dyn SeekableRead>, magnus::Error> {
+        if !Self::is_string_io(ruby, &input) {
+            return Err(magnus::Error::new(
+                magnus::exception::type_error(),
+                "Input is not a StringIO",
+            ));
+        }
+
+        let string_content = input.funcall::<_, _, RString>("string", ()).unwrap();
+        Ok(Box::new(Self {
+            inner: string_content,
+            offset: 0,
+        }))
+    }
+
+    fn is_string_io(ruby: &Ruby, input: &Value) -> bool {
+        let string_io_class = STRING_IO_CLASS.get_or_init(|| {
+            let class = RClass::from_value(ruby.eval("StringIO").unwrap()).unwrap();
+            Opaque::from(class)
+        });
+        input.is_kind_of(ruby.get_inner(*string_io_class))
+    }
+
+    fn from_string_like(input: Value) -> Result<Box<dyn SeekableRead>, magnus::Error> {
+        // Try calling `to_str`, and if that fails, try `to_s`
+        let string_content = input
+            .funcall::<_, _, RString>("to_str", ())
+            .or_else(|_| input.funcall::<_, _, RString>("to_s", ()))?;
+        Ok(Box::new(Self {
+            inner: string_content,
+            offset: 0,
+        }))
+    }
+}
+
+impl Read for RubyReader<Value> {
+    fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
+        let bytes = self
+            .inner
+            .funcall::<_, _, RString>("read", (buf.len(),))
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
+
+        buf.write_all(unsafe { bytes.as_slice() })?;
+
+        Ok(bytes.len())
+    }
+}
+
+impl Read for RubyReader<RString> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        let string_buffer = unsafe { self.inner.as_slice() };
+        if self.offset >= string_buffer.len() {
+            return Ok(0); // EOF
+        }
+
+        let remaining = string_buffer.len() - self.offset;
+        let copy_size = remaining.min(buf.len());
+        buf[..copy_size].copy_from_slice(&string_buffer[self.offset..self.offset + copy_size]);
+        self.offset += copy_size;
+        Ok(copy_size)
+    }
+}

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

@@ -0,0 +1,75 @@
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum ParserResultType {
+    Hash,
+    Array,
+}
+
+impl ParserResultType {
+    pub fn iter() -> impl Iterator<Item = Self> {
+        [Self::Hash, Self::Array].into_iter()
+    }
+}
+
+impl TryFrom<&str> for ParserResultType {
+    type Error = String;
+
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        match value {
+            "hash" => Ok(ParserResultType::Hash),
+            "array" => Ok(ParserResultType::Array),
+            _ => Err(format!("Invalid parser result type: {}", value)),
+        }
+    }
+}
+
+impl TryFrom<String> for ParserResultType {
+    type Error = String;
+
+    fn try_from(value: String) -> Result<Self, Self::Error> {
+        Self::try_from(value.as_str())
+    }
+}
+
+impl std::fmt::Display for ParserResultType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ParserResultType::Hash => write!(f, "hash"),
+            ParserResultType::Array => write!(f, "array"),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ListField<'a> {
+    pub item_type: ParquetSchemaType<'a>,
+    pub format: Option<&'a str>,
+}
+
+#[derive(Debug, Clone)]
+pub struct MapField<'a> {
+    pub key_type: ParquetSchemaType<'a>,
+    pub value_type: ParquetSchemaType<'a>,
+    pub format: Option<&'a str>,
+}
+
+#[derive(Debug, Clone)]
+pub enum ParquetSchemaType<'a> {
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    UInt8,
+    UInt16,
+    UInt32,
+    UInt64,
+    Float,
+    Double,
+    String,
+    Binary,
+    Boolean,
+    Date32,
+    TimestampMillis,
+    TimestampMicros,
+    List(Box<ListField<'a>>),
+    Map(Box<MapField<'a>>),
+}