parquet-tyfoom 0.8.0

data/ext/parquet-ruby-adapter/src/reader.rs

@@ -0,0 +1,339 @@
+use magnus::value::{BoxValue, ReprValue};
+use magnus::{Error as MagnusError, IntoValue, RArray, RHash, Ruby, TryConvert, Value};
+use parquet_core::reader::Reader;
+
+use crate::{
+    converter::parquet_to_ruby,
+    io::{RubyIOReader, ThreadSafeRubyIOReader},
+    logger::RubyLogger,
+    string_storage::{StringStorage, StringStorageConfig},
+    types::{ColumnEnumeratorArgs, ParserResultType, RowEnumeratorArgs},
+    utils::{create_column_enumerator, create_row_enumerator, handle_block_or_enum},
+    CloneableChunkReader,
+};
+use std::collections::HashSet;
+
+/// Read parquet file row by row
+pub fn each_row(
+    ruby: &Ruby,
+    rb_self: Value,
+    to_read: Value,
+    result_type: ParserResultType,
+    columns: Option<Vec<String>>,
+    strict: bool,
+    string_storage: StringStorageConfig,
+    logger: RubyLogger,
+) -> Result<Value, MagnusError> {
+    if let Some(enum_value) = handle_block_or_enum(ruby.block_given(), || {
+        create_row_enumerator(
+            ruby,
+            RowEnumeratorArgs {
+                rb_self,
+                to_read,
+                result_type,
+                columns: columns.clone(),
+                strict,
+                string_storage,
+                logger: logger.inner(),
+            },
+        )
+        .map(|yield_enum| yield_enum.into_value_with(ruby))
+    })? {
+        return Ok(enum_value);
+    }
+
+    // Log start of processing
+    let _ = logger.info(|| "Starting to read parquet file".to_string());
+
+    // Create a streaming reader based on input type
+    let chunk_reader = if to_read.is_kind_of(ruby.class_string()) {
+        let path_str: String = TryConvert::try_convert(to_read)?;
+        let _ = logger.debug(|| format!("Reading from file: {}", path_str));
+        CloneableChunkReader::from_path(&path_str)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?
+    } else if to_read.respond_to("read", false)? {
+        // Handle IO objects with streaming
+        let _ = logger.debug(|| "Reading from IO object".to_string());
+        let ruby_reader = RubyIOReader::new(to_read)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+        let thread_safe_reader = ThreadSafeRubyIOReader::new(ruby_reader);
+
+        CloneableChunkReader::from_ruby_io(thread_safe_reader)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?
+    } else {
+        return Err(MagnusError::new(
+            ruby.exception_runtime_error(),
+            format!(
+                "Invalid input type: expected String or IO object with read method, got {}",
+                to_read.class()
+            ),
+        ));
+    };
+
+    let reader = Reader::new(chunk_reader.clone());
+    let mut reader_for_metadata = Reader::new(chunk_reader);
+
+    // Get metadata to extract column names
+    let metadata = reader_for_metadata
+        .metadata()
+        .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+    let schema = metadata.schema();
+    let all_column_names: Vec<String> = schema
+        .get_fields()
+        .iter()
+        .map(|f| f.name().to_string())
+        .collect();
+
+    let _ = logger.info(|| format!("Processing {} columns", all_column_names.len()));
+
+    // Get the row iterator. Projected rows are yielded in file-schema order, not
+    // request order, so the hash keys must follow file order too — derive them by
+    // filtering the file columns, never from the request-ordered `cols`.
+    let (row_iter, column_names) = if let Some(ref cols) = columns {
+        let requested = cols.iter().map(String::as_str).collect::<HashSet<_>>();
+        let projected_names = all_column_names
+            .iter()
+            .filter(|name| requested.contains(name.as_str()))
+            .cloned()
+            .collect::<Vec<_>>();
+        let iter = reader
+            .read_rows_with_projection(cols)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+        (iter, projected_names)
+    } else {
+        let iter = reader
+            .read_rows()
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+        (iter, all_column_names)
+    };
+
+    let mut row_count = 0u64;
+    let mut string_storage = StringStorage::new(string_storage);
+
+    // BoxValue registers each interned key with the GC so it survives a
+    // GC.compact triggered by user code inside the yield loop; a bare RString
+    // held in this Vec would be relocated and dangle.
+    let interned_column_names = column_names
+        .iter()
+        .map(|name| BoxValue::new(ruby.str_new(name).to_interned_str()))
+        .collect::<Vec<_>>();
+
+    for row_result in row_iter {
+        let row = row_result
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+
+        // Convert row to Ruby value based on result_type
+        let ruby_row = match result_type {
+            ParserResultType::Array => {
+                let array: RArray = ruby.ary_new_capa(row.len());
+                for value in row {
+                    let ruby_value = parquet_to_ruby(value, &mut string_storage).map_err(|e| {
+                        MagnusError::new(ruby.exception_runtime_error(), e.to_string())
+                    })?;
+                    array.push(ruby_value)?;
+                }
+                array.as_value()
+            }
+            ParserResultType::Hash => {
+                let hash: RHash = ruby.hash_new_capa(row.len());
+                for (idx, value) in row.into_iter().enumerate() {
+                    if idx < interned_column_names.len() {
+                        let ruby_value =
+                            parquet_to_ruby(value, &mut string_storage).map_err(|e| {
+                                MagnusError::new(ruby.exception_runtime_error(), e.to_string())
+                            })?;
+                        hash.aset(interned_column_names[idx].as_value(), ruby_value)?;
+                    }
+                }
+                hash.as_value()
+            }
+        };
+
+        let _: Value = ruby.yield_value(ruby_row)?;
+
+        row_count += 1;
+        if row_count % 1000 == 0 {
+            let _ = logger.debug(|| format!("Processed {} rows", row_count));
+        }
+    }
+
+    let _ = logger.info(|| format!("Finished processing {} rows", row_count));
+
+    Ok(ruby.qnil().as_value())
+}
+
+/// Arguments for each_column function
+struct EachColumnArgs {
+    rb_self: Value,
+    to_read: Value,
+    result_type: ParserResultType,
+    columns: Option<Vec<String>>,
+    batch_size: Option<usize>,
+    strict: bool,
+    string_storage: StringStorageConfig,
+    logger: RubyLogger,
+}
+
+/// Read parquet file column by column
+#[allow(clippy::too_many_arguments)]
+pub fn each_column(
+    ruby: &Ruby,
+    rb_self: Value,
+    to_read: Value,
+    result_type: ParserResultType,
+    columns: Option<Vec<String>>,
+    batch_size: Option<usize>,
+    strict: bool,
+    string_storage: StringStorageConfig,
+    logger: RubyLogger,
+) -> Result<Value, MagnusError> {
+    let args = EachColumnArgs {
+        rb_self,
+        to_read,
+        result_type,
+        columns,
+        batch_size,
+        strict,
+        string_storage,
+        logger,
+    };
+    each_column_impl(ruby, args)
+}
+
+fn each_column_impl(ruby: &Ruby, args: EachColumnArgs) -> Result<Value, MagnusError> {
+    if let Some(enum_value) = handle_block_or_enum(ruby.block_given(), || {
+        create_column_enumerator(
+            ruby,
+            ColumnEnumeratorArgs {
+                rb_self: args.rb_self,
+                to_read: args.to_read,
+                result_type: args.result_type,
+                columns: args.columns.clone(),
+                batch_size: args.batch_size,
+                strict: args.strict,
+                string_storage: args.string_storage,
+                logger: args.logger.inner(),
+            },
+        )
+        .map(|yield_enum| yield_enum.into_value_with(ruby))
+    })? {
+        return Ok(enum_value);
+    }
+
+    // Log start of processing
+    let _ = args
+        .logger
+        .info(|| "Starting to read parquet file columns".to_string());
+
+    // Create a streaming reader based on input type
+    let chunk_reader = if args.to_read.is_kind_of(ruby.class_string()) {
+        let path_str: String = TryConvert::try_convert(args.to_read)?;
+        let _ = args
+            .logger
+            .debug(|| format!("Reading columns from file: {}", path_str));
+        CloneableChunkReader::from_path(&path_str)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?
+    } else if args.to_read.respond_to("read", false)? {
+        // Handle IO objects with streaming
+        let _ = args
+            .logger
+            .debug(|| "Reading columns from IO object".to_string());
+        let ruby_reader = RubyIOReader::new(args.to_read)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+        let thread_safe_reader = ThreadSafeRubyIOReader::new(ruby_reader);
+
+        CloneableChunkReader::from_ruby_io(thread_safe_reader)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?
+    } else {
+        return Err(MagnusError::new(
+            ruby.exception_runtime_error(),
+            format!(
+                "Invalid input type: expected String or IO object with read method, got {}",
+                args.to_read.class()
+            ),
+        ));
+    };
+
+    let reader = Reader::new(chunk_reader.clone());
+    let mut reader_for_metadata = Reader::new(chunk_reader);
+
+    // Get metadata to extract column names
+    let metadata = reader_for_metadata
+        .metadata()
+        .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+    let schema = metadata.schema();
+    let all_column_names: Vec<String> = schema
+        .get_fields()
+        .iter()
+        .map(|f| f.name().to_string())
+        .collect();
+
+    // Get the column iterator
+    let (col_iter, _column_names) = if let Some(ref cols) = args.columns {
+        let iter = reader
+            .read_columns_with_projection(cols, args.batch_size)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+        (iter, cols.clone())
+    } else {
+        let iter = reader
+            .read_columns(args.batch_size)
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+        (iter, all_column_names)
+    };
+
+    let mut batch_count = 0u64;
+    let mut string_storage = StringStorage::new(args.string_storage);
+
+    for batch_result in col_iter {
+        let batch = batch_result
+            .map_err(|e| MagnusError::new(ruby.exception_runtime_error(), e.to_string()))?;
+
+        // Convert batch to Ruby value based on result_type
+        let ruby_batch = match args.result_type {
+            ParserResultType::Array => {
+                let array: RArray = ruby.ary_new_capa(batch.columns.len());
+                for (_name, values) in batch.columns {
+                    let col_array: RArray = ruby.ary_new_capa(values.len());
+                    for value in values {
+                        let ruby_value =
+                            parquet_to_ruby(value, &mut string_storage).map_err(|e| {
+                                MagnusError::new(ruby.exception_runtime_error(), e.to_string())
+                            })?;
+                        col_array.push(ruby_value)?;
+                    }
+                    array.push(col_array)?;
+                }
+                array.as_value()
+            }
+            ParserResultType::Hash => {
+                let hash: RHash = ruby.hash_new();
+                for (name, values) in batch.columns {
+                    let col_array: RArray = ruby.ary_new_capa(values.len());
+                    for value in values {
+                        let ruby_value =
+                            parquet_to_ruby(value, &mut string_storage).map_err(|e| {
+                                MagnusError::new(ruby.exception_runtime_error(), e.to_string())
+                            })?;
+                        col_array.push(ruby_value)?;
+                    }
+                    let ruby_key = string_storage.ruby_key(ruby, &name);
+                    hash.aset(ruby_key, col_array)?;
+                }
+                hash.as_value()
+            }
+        };
+
+        let _: Value = ruby.yield_value(ruby_batch)?;
+
+        batch_count += 1;
+        let _ = args
+            .logger
+            .debug(|| format!("Processed batch {}", batch_count));
+    }
+
+    let _ = args
+        .logger
+        .info(|| format!("Finished processing {} batches", batch_count));
+
+    Ok(ruby.qnil().as_value())
+}