parquet 0.0.1 → 0.0.3

data/ext/parquet/src/reader.rs

@@ -1,122 +1,62 @@
-use crate::header_cache::{CacheError, StringCache};
-use crate::ruby_reader::{build_ruby_reader, SeekableRead};
-use crate::utils::*;
-use bytes::Bytes;
+// =============================================================================
+// Imports and Dependencies
+// =============================================================================
+use crate::header_cache::{CacheError, HeaderCacheCleanupIter, StringCache};
+use crate::{
+    create_column_enumerator, create_row_enumerator, utils::*, ColumnEnumeratorArgs, ColumnRecord,
+    ForgottenFileHandle, ParquetField, ParquetValueVec, RowEnumeratorArgs, RowRecord,
+    SeekableRubyValue,
+};
+use ahash::RandomState;
 use magnus::rb_sys::AsRawValue;
 use magnus::value::{Opaque, ReprValue};
-use magnus::IntoValue;
-use magnus::{block::Yield, Error as MagnusError, KwArgs, RHash, Ruby, Symbol, Value};
+use magnus::{block::Yield, Error as MagnusError, Ruby, Value};
+use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
+use parquet::arrow::ProjectionMask;
 use parquet::errors::ParquetError;
-use parquet::file::reader::{ChunkReader, Length, SerializedFileReader};
-use parquet::record::Field;
+use parquet::file::reader::FileReader;
+use parquet::file::reader::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::io::{self, BufReader, Read, Seek, SeekFrom};
+use std::io::{self};
 use std::mem::ManuallyDrop;
 use std::os::fd::FromRawFd;
 use std::sync::OnceLock;
-use std::{borrow::Cow, hash::BuildHasher};
 use thiserror::Error;
-use xxhash_rust::xxh3::Xxh3Builder;
-
-use parquet::record::reader::RowIter as ParquetRowIter;
-
-#[derive(Error, Debug)]
-pub enum ReaderError {
-    #[error("Failed to get file descriptor: {0}")]
-    FileDescriptor(String),
-    #[error("Invalid file descriptor")]
-    InvalidFileDescriptor,
-    #[error("Failed to open file: {0}")]
-    FileOpen(#[from] io::Error),
-    #[error("Failed to intern headers: {0}")]
-    HeaderIntern(#[from] CacheError),
-    #[error("Ruby error: {0}")]
-    Ruby(String),
-}
-
-impl From<MagnusError> for ReaderError {
-    fn from(err: MagnusError) -> Self {
-        Self::Ruby(err.to_string())
-    }
-}
-
-impl From<ReaderError> for MagnusError {
-    fn from(err: ReaderError) -> Self {
-        MagnusError::new(
-            Ruby::get().unwrap().exception_runtime_error(),
-            err.to_string(),
-        )
-    }
-}
-
-struct ForgottenFileHandle(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)
-    }
-}
-
-struct HeaderCacheCleanupIter<I> {
-    inner: I,
-    headers: OnceLock<Vec<&'static str>>,
-}
-
-impl<I: Iterator> Iterator for HeaderCacheCleanupIter<I> {
-    type Item = I::Item;
 
-    fn next(&mut self) -> Option<Self::Item> {
-        self.inner.next()
-    }
-}
-
-impl<I> Drop for HeaderCacheCleanupIter<I> {
-    fn drop(&mut self) {
-        if let Some(headers) = self.headers.get() {
-            StringCache::clear(&headers).unwrap();
-        }
-    }
-}
-
-pub fn parse_parquet<'a>(
+#[inline]
+pub fn parse_parquet_rows<'a>(
     rb_self: Value,
     args: &[Value],
-) -> Result<Yield<Box<dyn Iterator<Item = Record<Xxh3Builder>>>>, MagnusError> {
+) -> Result<Yield<Box<dyn Iterator<Item = RowRecord<RandomState>>>>, MagnusError> {
     let original = unsafe { Ruby::get_unchecked() };
     let ruby: &'static Ruby = Box::leak(Box::new(original));
 
-    let ParquetArgs {
+    let ParquetRowsArgs {
         to_read,
         result_type,
-    } = parse_parquet_args(&ruby, args)?;
+        columns,
+    } = parse_parquet_rows_args(&ruby, args)?;
 
     if !ruby.block_given() {
-        return create_enumerator(EnumeratorArgs {
+        return create_row_enumerator(RowEnumeratorArgs {
             rb_self,
             to_read,
             result_type,
+            columns,
         });
     }
 
-    let iter = if to_read.is_kind_of(ruby.class_string()) {
+    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();
-        ParquetRowIter::from_file_into(Box::new(reader))
+        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) })
@@ -131,14 +71,28 @@ pub fn parse_parquet<'a>(
         let file = unsafe { File::from_raw_fd(fd) };
         let file = ForgottenFileHandle(ManuallyDrop::new(file));
         let reader = SerializedFileReader::new(file).unwrap();
-        ParquetRowIter::from_file_into(Box::new(reader))
+        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();
-        ParquetRowIter::from_file_into(Box::new(reader))
+        let schema = reader.metadata().file_metadata().schema().clone();
+
+        (schema, ParquetRowIter::from_file_into(Box::new(reader)))
     };
 
-    let iter: Box<dyn Iterator<Item = Record<Xxh3Builder>>> = match result_type.as_str() {
+    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),
+            )
+        })?;
+    }
+
+    let iter: Box<dyn Iterator<Item = RowRecord<RandomState>>> = match result_type.as_str() {
         "hash" => {
             let headers = OnceLock::new();
             let headers_clone = headers.clone();
@@ -146,21 +100,27 @@ pub fn parse_parquet<'a>(
                 .filter_map(move |row| {
                     row.ok().map(|row| {
                         let headers = headers_clone.get_or_init(|| {
-                            row.get_column_iter()
-                                .map(|(k, _)| StringCache::intern(k.to_owned()).unwrap())
-                                .collect::<Vec<_>>()
+                            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
                         });
 
-                        row.get_column_iter()
-                            .enumerate()
-                            .map(|(i, (_, v))| {
-                                let key = headers[i];
-                                (key, ParquetField(v.clone()))
-                            })
-                            .collect::<HashMap<&'static str, ParquetField, Xxh3Builder>>()
+                        let mut map =
+                            HashMap::with_capacity_and_hasher(headers.len(), Default::default());
+                        row.get_column_iter().enumerate().for_each(|(i, (_, v))| {
+                            map.insert(headers[i], ParquetField(v.clone()));
+                        });
+                        map
                     })
                 })
-                .map(|row| Record::Map(row));
+                .map(RowRecord::Map);
 
             Box::new(HeaderCacheCleanupIter {
                 inner: iter,
@@ -170,12 +130,14 @@ pub fn parse_parquet<'a>(
         "array" => Box::new(
             iter.filter_map(|row| {
                 row.ok().map(|row| {
+                    let column_count = row.get_column_iter().count();
+                    let mut vec = Vec::with_capacity(column_count);
                     row.get_column_iter()
-                        .map(|(_, v)| ParquetField(v.clone()))
-                        .collect::<Vec<ParquetField>>()
+                        .for_each(|(_, v)| vec.push(ParquetField(v.clone())));
+                    vec
                 })
             })
-            .map(|row| Record::Vec(row)),
+            .map(RowRecord::Vec),
         ),
         _ => {
             return Err(MagnusError::new(
@@ -188,150 +150,246 @@ pub fn parse_parquet<'a>(
     Ok(Yield::Iter(iter))
 }
 
-struct EnumeratorArgs {
+#[inline]
+pub fn parse_parquet_columns<'a>(
     rb_self: Value,
-    to_read: Value,
-    result_type: String,
-}
+    args: &[Value],
+) -> Result<Yield<Box<dyn Iterator<Item = ColumnRecord<RandomState>>>>, MagnusError> {
+    let original = unsafe { Ruby::get_unchecked() };
+    let ruby: &'static Ruby = Box::leak(Box::new(original));
 
-fn create_enumerator(
-    args: EnumeratorArgs,
-) -> Result<Yield<Box<dyn Iterator<Item = Record<Xxh3Builder>>>>, MagnusError> {
-    let kwargs = RHash::new();
+    let ParquetColumnsArgs {
+        to_read,
+        result_type,
+        columns,
+        batch_size,
+    } = parse_parquet_columns_args(&ruby, args)?;
 
-    kwargs.aset(Symbol::new("result_type"), Symbol::new(args.result_type))?;
+    if !ruby.block_given() {
+        return create_column_enumerator(ColumnEnumeratorArgs {
+            rb_self,
+            to_read,
+            result_type,
+            columns,
+            batch_size,
+        });
+    }
 
-    let enumerator = args
-        .rb_self
-        .enumeratorize("for_each", (args.to_read, KwArgs(kwargs)));
-    Ok(Yield::Enumerator(enumerator))
-}
+    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))?;
 
-#[derive(Debug)]
-pub enum Record<S: BuildHasher + Default> {
-    Vec(Vec<ParquetField>),
-    Map(HashMap<&'static str, ParquetField, S>),
-}
+        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();
 
-impl<S: BuildHasher + Default> IntoValue for Record<S> {
-    #[inline]
-    fn into_value_with(self, handle: &Ruby) -> Value {
-        match self {
-            Record::Vec(vec) => {
-                let ary = handle.ary_new_capa(vec.len());
-                vec.into_iter().try_for_each(|v| ary.push(v)).unwrap();
-                ary.into_value_with(handle)
-            }
-            Record::Map(map) => {
-                // Pre-allocate the hash with the known size
-                let hash = handle.hash_new_capa(map.len());
-                map.into_iter()
-                    .try_for_each(|(k, v)| hash.aset(k, v))
-                    .unwrap();
-                hash.into_value_with(handle)
-            }
+        // 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);
         }
-    }
-}
 
-#[derive(Debug, Clone)]
-pub struct CowValue<'a>(pub Cow<'a, str>);
+        let reader = builder.build().unwrap();
 
-impl<'a> IntoValue for CowValue<'a> {
-    fn into_value_with(self, handle: &Ruby) -> Value {
-        self.0.into_value_with(handle)
+        (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).unwrap();
+        let schema = builder.schema().clone();
+        let num_rows = builder.metadata().file_metadata().num_rows();
+
+        if let Some(batch_size) = batch_size {
+            builder = builder.with_batch_size(batch_size);
+        }
+
+        // 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);
+        }
+
+        let reader = builder.build().unwrap();
+
+        (reader, schema, num_rows)
+    } else {
+        let readable = SeekableRubyValue(Opaque::from(to_read));
+
+        let mut builder = ParquetRecordBatchReaderBuilder::try_new(readable).unwrap();
+        let schema = builder.schema().clone();
+        let num_rows = builder.metadata().file_metadata().num_rows();
+
+        if let Some(batch_size) = batch_size {
+            builder = builder.with_batch_size(batch_size);
+        }
+
+        // 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);
+        }
+
+        let reader = builder.build().unwrap();
+
+        (reader, schema, num_rows)
+    };
+
+    if num_rows == 0 {
+        let mut map =
+            HashMap::with_capacity_and_hasher(schema.fields().len(), RandomState::default());
+        for field in schema.fields() {
+            map.insert(
+                StringCache::intern(field.name().to_string()).unwrap(),
+                vec![],
+            );
+        }
+        let column_record = vec![ColumnRecord::Map(map)];
+        return Ok(Yield::Iter(Box::new(column_record.into_iter())));
     }
-}
 
-#[derive(Debug)]
-pub struct ParquetField(Field);
-
-impl<'a> IntoValue for ParquetField {
-    fn into_value_with(self, handle: &Ruby) -> Value {
-        match self.0 {
-            Field::Byte(b) => b.into_value_with(handle),
-            Field::Bool(b) => b.into_value_with(handle),
-            Field::Short(s) => s.into_value_with(handle),
-            Field::Int(i) => i.into_value_with(handle),
-            Field::Long(l) => l.into_value_with(handle),
-            Field::UByte(ub) => ub.into_value_with(handle),
-            Field::UShort(us) => us.into_value_with(handle),
-            Field::UInt(ui) => ui.into_value_with(handle),
-            Field::ULong(ul) => ul.into_value_with(handle),
-            Field::Float16(f) => f32::from(f).into_value_with(handle),
-            Field::Float(f) => f.into_value_with(handle),
-            Field::Double(d) => d.into_value_with(handle),
-
-            Field::Str(s) => s.into_value_with(handle),
-            Field::Bytes(b) => handle.str_from_slice(b.data()).as_value(),
-            Field::Date(d) => d.into_value_with(handle),
-            Field::TimestampMillis(ts) => ts.into_value_with(handle),
-            Field::TimestampMicros(ts) => ts.into_value_with(handle),
-            Field::ListInternal(list) => {
-                let ary = handle.ary_new_capa(list.elements().len());
-                list.elements()
-                    .iter()
-                    .try_for_each(|e| ary.push(ParquetField(e.clone()).into_value_with(handle)))
-                    .unwrap();
-                ary.into_value_with(handle)
-            }
-            Field::MapInternal(map) => {
-                let hash = handle.hash_new_capa(map.entries().len());
-                map.entries()
-                    .iter()
-                    .try_for_each(|(k, v)| {
-                        hash.aset(
-                            ParquetField(k.clone()).into_value_with(handle),
-                            ParquetField(v.clone()).into_value_with(handle),
-                        )
+    let iter: Box<dyn Iterator<Item = ColumnRecord<RandomState>>> = match result_type.as_str() {
+        "hash" => {
+            let headers = OnceLock::new();
+            let headers_clone = headers.clone();
+            let iter = batch_reader
+                .filter_map(move |batch| {
+                    batch.ok().map(|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(), Default::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());
+                        });
+
+                        map
                     })
-                    .unwrap();
-                hash.into_value_with(handle)
-            }
-            // Field::Decimal(d) => d.to_string().into_value_with(handle),
-            // Field::Group(row) => row.into_value_with(handle),
-            Field::Null => handle.qnil().as_value(),
-            _ => panic!("Unsupported field type"),
+                })
+                .map(ColumnRecord::Map);
+
+            Box::new(HeaderCacheCleanupIter {
+                inner: iter,
+                headers,
+            })
         }
-    }
-}
+        "array" => Box::new(
+            batch_reader
+                .filter_map(|batch| {
+                    batch.ok().map(|batch| {
+                        batch
+                            .columns()
+                            .into_iter()
+                            .map(|column| {
+                                let values = ParquetValueVec::try_from(column.clone()).unwrap();
+                                values.into_inner()
+                            })
+                            .collect()
+                    })
+                })
+                .map(ColumnRecord::Vec),
+        ),
+        _ => {
+            return Err(MagnusError::new(
+                ruby.exception_runtime_error(),
+                "Invalid result type",
+            ))
+        }
+    };
 
-struct SeekableRubyValue(Opaque<Value>);
+    Ok(Yield::Iter(iter))
+}
 
-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
+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()
     }
 }
 
-impl ChunkReader for SeekableRubyValue {
-    type T = BufReader<Box<dyn SeekableRead>>;
+#[derive(Error, Debug)]
+pub enum ReaderError {
+    #[error("Failed to get file descriptor: {0}")]
+    FileDescriptor(String),
+    #[error("Invalid file descriptor")]
+    InvalidFileDescriptor,
+    #[error("Failed to open file: {0}")]
+    FileOpen(#[from] io::Error),
+    #[error("Failed to intern headers: {0}")]
+    HeaderIntern(#[from] CacheError),
+    #[error("Ruby error: {0}")]
+    Ruby(String),
+    #[error("Parquet error: {0}")]
+    Parquet(#[from] ParquetError),
+}
 
-    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::new(reader))
+impl From<MagnusError> for ReaderError {
+    fn from(err: MagnusError) -> Self {
+        Self::Ruby(err.to_string())
     }
+}
 
-    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())
+impl From<ReaderError> for MagnusError {
+    fn from(err: ReaderError) -> Self {
+        MagnusError::new(
+            Ruby::get().unwrap().exception_runtime_error(),
+            err.to_string(),
+        )
     }
 }