parquet 0.4.2 → 0.5.1

data/ext/parquet/src/ruby_reader.rs

@@ -13,14 +13,18 @@ use std::{
     sync::Arc,
 };
 
+use crate::types::ParquetGemError;
+
 /// A reader that can handle various Ruby input types (String, StringIO, IO-like objects)
 /// and provide a standard Read implementation for them.
 pub enum RubyReader {
     String {
+        ruby: Arc<Ruby>,
         inner: Opaque<RString>,
         offset: usize,
     },
     RubyIoLike {
+        ruby: Arc<Ruby>,
         inner: Opaque<Value>,
     },
     NativeProxyIoLike {
@@ -28,26 +32,15 @@ pub enum RubyReader {
     },
 }
 
-impl RubyReader {
-    fn is_io_like(value: &Value) -> bool {
-        value.respond_to("read", false).unwrap_or(false)
-    }
-
-    // For now, don't use this. Having to use seek in length is scary.
-    fn is_seekable_io_like(value: &Value) -> bool {
-        Self::is_io_like(value)
-            && value.respond_to("seek", false).unwrap_or(false)
-            && value.respond_to("pos", false).unwrap_or(false)
-    }
-}
-
-impl TryFrom<Value> for RubyReader {
-    type Error = magnus::Error;
+// Sending is technically not safe, but the only things that threatens to
+// do this is the parquet gem, and they don't seem to actually do it.
+unsafe impl Send for RubyReader {}
 
-    fn try_from(value: Value) -> Result<Self, Self::Error> {
-        let ruby = unsafe { Ruby::get_unchecked() };
+impl RubyReader {
+    pub fn new(ruby: Arc<Ruby>, value: Value) -> Result<Self, ParquetGemError> {
         if RubyReader::is_seekable_io_like(&value) {
             Ok(RubyReader::RubyIoLike {
+                ruby,
                 inner: Opaque::from(value),
             })
         } else if RubyReader::is_io_like(&value) {
@@ -56,6 +49,7 @@ impl TryFrom<Value> for RubyReader {
 
             // This is safe, because we won't call seek
             let inner_readable = RubyReader::RubyIoLike {
+                ruby: ruby.clone(),
                 inner: Opaque::from(value),
             };
             let mut reader = BufReader::new(inner_readable);
@@ -74,19 +68,31 @@ impl TryFrom<Value> for RubyReader {
                 .funcall::<_, _, RString>("to_str", ())
                 .or_else(|_| value.funcall::<_, _, RString>("to_s", ()))?;
             Ok(RubyReader::String {
+                ruby,
                 inner: Opaque::from(string_content),
                 offset: 0,
             })
         }
     }
+
+    fn is_io_like(value: &Value) -> bool {
+        value.respond_to("read", false).unwrap_or(false)
+    }
+
+    // For now, don't use this. Having to use seek in length is scary.
+    fn is_seekable_io_like(value: &Value) -> bool {
+        Self::is_io_like(value)
+            && value.respond_to("seek", false).unwrap_or(false)
+            && value.respond_to("pos", false).unwrap_or(false)
+    }
 }
 
 impl Seek for RubyReader {
     fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
-        let ruby = unsafe { Ruby::get_unchecked() };
         match self {
             RubyReader::NativeProxyIoLike { proxy_file } => proxy_file.seek(pos),
             RubyReader::String {
+                ruby,
                 inner,
                 offset: original_offset,
             } => {
@@ -107,7 +113,7 @@ impl Seek for RubyReader {
                 *original_offset = new_offset.min(unwrapped_inner.len());
                 Ok(*original_offset as u64)
             }
-            RubyReader::RubyIoLike { inner } => {
+            RubyReader::RubyIoLike { ruby, inner } => {
                 let unwrapped_inner = ruby.get_inner(*inner);
 
                 let (whence, ruby_offset) = match pos {
@@ -132,10 +138,13 @@ impl Seek for RubyReader {
 
 impl Read for RubyReader {
     fn read(&mut self, mut buf: &mut [u8]) -> io::Result<usize> {
-        let ruby = unsafe { Ruby::get_unchecked() };
         match self {
             RubyReader::NativeProxyIoLike { proxy_file } => proxy_file.read(buf),
-            RubyReader::String { inner, offset } => {
+            RubyReader::String {
+                ruby,
+                inner,
+                offset,
+            } => {
                 let unwrapped_inner = ruby.get_inner(*inner);
 
                 let string_buffer = unsafe { unwrapped_inner.as_slice() };
@@ -151,7 +160,7 @@ impl Read for RubyReader {
 
                 Ok(copy_size)
             }
-            RubyReader::RubyIoLike { inner } => {
+            RubyReader::RubyIoLike { ruby, inner } => {
                 let unwrapped_inner = ruby.get_inner(*inner);
 
                 let bytes = unwrapped_inner
@@ -175,14 +184,17 @@ impl Read for RubyReader {
 
 impl Length for RubyReader {
     fn len(&self) -> u64 {
-        let ruby = unsafe { Ruby::get_unchecked() };
         match self {
             RubyReader::NativeProxyIoLike { proxy_file } => proxy_file.len(),
-            RubyReader::String { inner, offset: _ } => {
+            RubyReader::String {
+                ruby,
+                inner,
+                offset: _,
+            } => {
                 let unwrapped_inner = ruby.get_inner(*inner);
                 unwrapped_inner.len() as u64
             }
-            RubyReader::RubyIoLike { inner } => {
+            RubyReader::RubyIoLike { ruby, inner } => {
                 let unwrapped_inner = ruby.get_inner(*inner);
 
                 // Get current position

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

@@ -43,17 +43,60 @@ impl std::fmt::Display for ParserResultType {
 pub struct ListField<'a> {
     pub item_type: ParquetSchemaType<'a>,
     pub format: Option<&'a str>,
+    pub nullable: bool,
 }
 
 #[derive(Debug, Clone)]
 pub struct MapField<'a> {
     pub key_type: ParquetSchemaType<'a>,
     pub value_type: ParquetSchemaType<'a>,
-    pub format: Option<&'a str>,
+    pub key_format: Option<&'a str>,
+    pub value_format: Option<&'a str>,
+    pub value_nullable: bool,
 }
 
 #[derive(Debug, Clone)]
+pub struct StructField<'a> {
+    pub fields: Vec<super::writer_types::SchemaField<'a>>,
+}
+
+#[derive(Clone, Debug)]
 pub enum ParquetSchemaType<'a> {
+    Primitive(PrimitiveType),
+    List(Box<ListField<'a>>),
+    Map(Box<MapField<'a>>),
+    Struct(Box<StructField<'a>>),
+}
+
+// New schema representation for the DSL-based approach
+#[derive(Debug, Clone)]
+pub enum SchemaNode {
+    Struct {
+        name: String,
+        nullable: bool,
+        fields: Vec<SchemaNode>,
+    },
+    List {
+        name: String,
+        nullable: bool,
+        item: Box<SchemaNode>,
+    },
+    Map {
+        name: String,
+        nullable: bool,
+        key: Box<SchemaNode>,
+        value: Box<SchemaNode>,
+    },
+    Primitive {
+        name: String,
+        parquet_type: PrimitiveType,
+        nullable: bool,
+        format: Option<String>,
+    },
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum PrimitiveType {
     Int8,
     Int16,
     Int32,
@@ -62,14 +105,12 @@ pub enum ParquetSchemaType<'a> {
     UInt16,
     UInt32,
     UInt64,
-    Float,
-    Double,
+    Float32,
+    Float64,
+    Boolean,
     String,
     Binary,
-    Boolean,
     Date32,
     TimestampMillis,
     TimestampMicros,
-    List(Box<ListField<'a>>),
-    Map(Box<MapField<'a>>),
 }

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

@@ -2,13 +2,20 @@
 mod core_types;
 mod parquet_value;
 mod record_types;
+pub mod schema_converter;
+pub mod schema_node;
 mod timestamp;
-mod type_conversion;
+pub mod type_conversion;
 mod writer_types;
 
 pub use core_types::*;
 pub use parquet_value::*;
 pub use record_types::*;
+// Explicitly export schema-related items
+pub use schema_converter::{
+    infer_schema_from_first_row, legacy_schema_to_dsl, parse_legacy_schema,
+};
+pub use schema_node::parse_schema_node;
 pub use timestamp::*;
 pub use type_conversion::*;
 pub use writer_types::*;
@@ -28,3 +35,59 @@ use parquet::record::Field;
 use std::{collections::HashMap, hash::BuildHasher, sync::Arc};
 
 use crate::header_cache::StringCacheKey;
+
+use crate::header_cache::CacheError;
+
+use std::io;
+
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum ParquetGemError {
+    #[error("Failed to open file: {0}")]
+    FileOpen(#[from] io::Error),
+    #[error("Failed to intern headers: {0}")]
+    HeaderIntern(#[from] CacheError),
+    #[error("Ruby error: {0}")]
+    Ruby(#[from] MagnusErrorWrapper),
+    #[error("Parquet error: {0}")]
+    Parquet(#[from] parquet::errors::ParquetError),
+    #[error("Arrow error: {0}")]
+    Arrow(#[from] arrow_schema::ArrowError),
+    #[error("UTF-8 error: {0}")]
+    Utf8Error(#[from] simdutf8::basic::Utf8Error),
+    #[error("Jiff error: {0}")]
+    Jiff(#[from] jiff::Error),
+}
+
+#[derive(Debug)]
+pub struct MagnusErrorWrapper(pub MagnusError);
+
+impl From<MagnusError> for MagnusErrorWrapper {
+    fn from(err: MagnusError) -> Self {
+        Self(err)
+    }
+}
+
+impl std::fmt::Display for MagnusErrorWrapper {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl std::error::Error for MagnusErrorWrapper {}
+
+impl From<MagnusError> for ParquetGemError {
+    fn from(err: MagnusError) -> Self {
+        Self::Ruby(MagnusErrorWrapper(err))
+    }
+}
+
+impl Into<MagnusError> for ParquetGemError {
+    fn into(self) -> MagnusError {
+        match self {
+            Self::Ruby(MagnusErrorWrapper(err)) => err.into(),
+            _ => MagnusError::new(magnus::exception::runtime_error(), self.to_string()),
+        }
+    }
+}