parquet 0.4.2 → 0.5.0

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

@@ -0,0 +1,349 @@
+use magnus::value::ReprValue; // Add ReprValue trait to scope
+use magnus::{Error as MagnusError, RArray, Ruby, TryConvert, Value};
+
+use crate::types::{ParquetSchemaType as PST, PrimitiveType, SchemaField, SchemaNode};
+use crate::utils::parse_string_or_symbol;
+
+/// Recursively converts a SchemaField to a SchemaNode for any level of nesting
+fn convert_schema_field_to_node(field: &SchemaField) -> SchemaNode {
+    match &field.type_ {
+        PST::Int8 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Int8,
+            format: field.format.clone(),
+        },
+        PST::Int16 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Int16,
+            format: field.format.clone(),
+        },
+        PST::Int32 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Int32,
+            format: field.format.clone(),
+        },
+        PST::Int64 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Int64,
+            format: field.format.clone(),
+        },
+        PST::UInt8 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::UInt8,
+            format: field.format.clone(),
+        },
+        PST::UInt16 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::UInt16,
+            format: field.format.clone(),
+        },
+        PST::UInt32 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::UInt32,
+            format: field.format.clone(),
+        },
+        PST::UInt64 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::UInt64,
+            format: field.format.clone(),
+        },
+        PST::Float => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Float32,
+            format: field.format.clone(),
+        },
+        PST::Double => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Float64,
+            format: field.format.clone(),
+        },
+        PST::String => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::String,
+            format: field.format.clone(),
+        },
+        PST::Binary => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Binary,
+            format: field.format.clone(),
+        },
+        PST::Boolean => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Boolean,
+            format: field.format.clone(),
+        },
+        PST::Date32 => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::Date32,
+            format: field.format.clone(),
+        },
+        PST::TimestampMillis => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::TimestampMillis,
+            format: field.format.clone(),
+        },
+        PST::TimestampMicros => SchemaNode::Primitive {
+            name: field.name.clone(),
+            nullable: field.nullable,
+            parquet_type: PrimitiveType::TimestampMicros,
+            format: field.format.clone(),
+        },
+        PST::List(list_field) => {
+            // Create item node by recursively converting the list item type to a node
+            let item_node = match &list_field.item_type {
+                // For primitive types, create a primitive node with name "item"
+                PST::Int8
+                | PST::Int16
+                | PST::Int32
+                | PST::Int64
+                | PST::UInt8
+                | PST::UInt16
+                | PST::UInt32
+                | PST::UInt64
+                | PST::Float
+                | PST::Double
+                | PST::String
+                | PST::Binary
+                | PST::Boolean
+                | PST::Date32
+                | PST::TimestampMillis
+                | PST::TimestampMicros => {
+                    // Use a temporary SchemaField to convert item type
+                    let item_field = SchemaField {
+                        name: "item".to_string(),
+                        type_: list_field.item_type.clone(),
+                        format: list_field.format.clone().map(String::from),
+                        nullable: list_field.nullable,
+                    };
+                    convert_schema_field_to_node(&item_field)
+                }
+                // For nested types (List, Map, Struct), recursively convert them
+                PST::List(_) | PST::Map(_) | PST::Struct(_) => {
+                    // Use a temporary SchemaField to convert item type
+                    let item_field = SchemaField {
+                        name: "item".to_string(),
+                        type_: list_field.item_type.clone(),
+                        format: list_field.format.clone().map(String::from),
+                        nullable: list_field.nullable,
+                    };
+                    convert_schema_field_to_node(&item_field)
+                }
+            };
+
+            SchemaNode::List {
+                name: field.name.clone(),
+                nullable: field.nullable,
+                item: Box::new(item_node),
+            }
+        }
+        PST::Map(map_field) => {
+            let key_field = SchemaField {
+                name: "key".to_string(),
+                type_: map_field.key_type.clone(),
+                format: map_field.key_format.clone().map(String::from),
+                nullable: false, // Map keys can never be null in Parquet
+            };
+            let value_field = SchemaField {
+                name: "value".to_string(),
+                type_: map_field.value_type.clone(),
+                format: map_field.value_format.clone().map(String::from),
+                nullable: map_field.value_nullable,
+            };
+
+            let key_node = convert_schema_field_to_node(&key_field);
+            let value_node = convert_schema_field_to_node(&value_field);
+
+            SchemaNode::Map {
+                name: field.name.clone(),
+                nullable: field.nullable,
+                key: Box::new(key_node),
+                value: Box::new(value_node),
+            }
+        }
+        PST::Struct(struct_field) => {
+            // Convert each subfield recursively
+            let mut field_nodes = Vec::with_capacity(struct_field.fields.len());
+
+            for subfield in struct_field.fields.iter() {
+                // Recursively convert each subfield, supporting any level of nesting
+                field_nodes.push(convert_schema_field_to_node(subfield));
+            }
+
+            SchemaNode::Struct {
+                name: field.name.clone(),
+                nullable: field.nullable,
+                fields: field_nodes,
+            }
+        }
+    }
+}
+
+/// Converts the legacy schema format (array of field hashes) to the new DSL format (SchemaNode)
+pub fn legacy_schema_to_dsl(
+    _ruby: &Ruby,
+    schema_fields: Vec<SchemaField>,
+) -> Result<SchemaNode, MagnusError> {
+    // Create a top-level struct node with fields for each schema field
+    let mut field_nodes = Vec::with_capacity(schema_fields.len());
+
+    for field in schema_fields {
+        // Use our recursive converter to handle any level of nesting
+        field_nodes.push(convert_schema_field_to_node(&field));
+    }
+
+    Ok(SchemaNode::Struct {
+        name: "".to_string(), // Top level has no name
+        nullable: false,      // Top level is not nullable
+        fields: field_nodes,
+    })
+}
+
+/// Parses the legacy format schema (array of field hashes)
+pub fn parse_legacy_schema(
+    ruby: &Ruby,
+    schema_value: Value,
+) -> Result<Vec<SchemaField>, MagnusError> {
+    if schema_value.is_nil()
+        || (schema_value.is_kind_of(ruby.class_array())
+            && RArray::from_value(schema_value)
+                .ok_or_else(|| {
+                    MagnusError::new(
+                        ruby.exception_type_error(),
+                        "Schema must be an array of field definitions or nil",
+                    )
+                })?
+                .len()
+                == 0)
+    {
+        // If schema is nil or an empty array, we'll handle this in the caller
+        return Ok(Vec::new());
+    }
+
+    if schema_value.is_kind_of(ruby.class_array()) {
+        let schema_array = RArray::from_value(schema_value).ok_or_else(|| {
+            MagnusError::new(
+                ruby.exception_type_error(),
+                "Schema must be an array of field definitions or nil",
+            )
+        })?;
+        let mut schema = Vec::with_capacity(schema_array.len());
+
+        for (idx, field_hash) in schema_array.into_iter().enumerate() {
+            if !field_hash.is_kind_of(ruby.class_hash()) {
+                return Err(MagnusError::new(
+                    ruby.exception_type_error(),
+                    format!("schema[{}] must be a hash", idx),
+                ));
+            }
+
+            let entries: Vec<(Value, Value)> = field_hash.funcall("to_a", ())?;
+            if entries.len() != 1 {
+                return Err(MagnusError::new(
+                    ruby.exception_type_error(),
+                    format!("schema[{}] must contain exactly one key-value pair", idx),
+                ));
+            }
+
+            let (name, type_value) = &entries[0];
+            let name_option = parse_string_or_symbol(ruby, name.clone())?;
+            let name = name_option.ok_or_else(|| {
+                MagnusError::new(ruby.exception_runtime_error(), "Field name cannot be nil")
+            })?;
+
+            let (type_, format, nullable) = if type_value.is_kind_of(ruby.class_hash()) {
+                let type_hash: Vec<(Value, Value)> = type_value.funcall("to_a", ())?;
+                let mut type_str = None;
+                let mut format_str = None;
+                let mut nullable = true; // Default to true if not specified
+
+                for (key, value) in type_hash {
+                    let key_option = parse_string_or_symbol(ruby, key)?;
+                    let key = key_option.ok_or_else(|| {
+                        MagnusError::new(ruby.exception_runtime_error(), "Type key cannot be nil")
+                    })?;
+                    match key.as_str() {
+                        "type" => type_str = Some(value),
+                        "format" => {
+                            let format_option = parse_string_or_symbol(ruby, value)?;
+                            format_str = format_option;
+                        }
+                        "nullable" => {
+                            // Extract nullable if present - convert to boolean
+                            nullable = bool::try_convert(value).unwrap_or(true);
+                        }
+                        _ => {
+                            return Err(MagnusError::new(
+                                ruby.exception_type_error(),
+                                format!("Unknown key '{}' in type definition", key),
+                            ))
+                        }
+                    }
+                }
+
+                let type_str = type_str.ok_or_else(|| {
+                    MagnusError::new(
+                        ruby.exception_type_error(),
+                        "Missing 'type' in type definition",
+                    )
+                })?;
+
+                (PST::try_convert(type_str)?, format_str, nullable)
+            } else {
+                (PST::try_convert(type_value.clone())?, None, true)
+            };
+
+            schema.push(SchemaField {
+                name,
+                type_,
+                format,
+                nullable,
+            });
+        }
+
+        Ok(schema)
+    } else {
+        Err(MagnusError::new(
+            ruby.exception_type_error(),
+            "Schema must be an array of field definitions or nil",
+        ))
+    }
+}
+
+/// Generates schema fields by inferring from the first row
+pub fn infer_schema_from_first_row(
+    ruby: &Ruby,
+    first_value: Value,
+    nullable: bool,
+) -> Result<Vec<SchemaField>, MagnusError> {
+    let array = RArray::from_value(first_value).ok_or_else(|| {
+        MagnusError::new(
+            ruby.exception_type_error(),
+            "First value must be an array when schema is not provided",
+        )
+    })?;
+
+    // Generate field names f0, f1, f2, etc.
+    Ok((0..array.len())
+        .map(|i| SchemaField {
+            name: format!("f{}", i),
+            type_: PST::String, // Default to String type when inferring
+            format: None,
+            nullable,
+        })
+        .collect())
+}

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

@@ -0,0 +1,329 @@
+use std::sync::Arc;
+
+use arrow_schema::{
+    DataType as ArrowDataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema,
+};
+use magnus::{Error as MagnusError, RArray, RHash, Ruby, Symbol, TryConvert, Value};
+
+use crate::logger::RubyLogger;
+use crate::types::{PrimitiveType, SchemaNode};
+use crate::utils::parse_string_or_symbol;
+
+/// Builds an Arrow schema from a SchemaNode tree - placeholder declaration
+/// The actual implementation appears later in the file
+fn _build_arrow_schema_placeholder() {}
+
+/// Helper to extract common fields from a schema node hash
+fn extract_common_fields(
+    ruby: &Ruby,
+    node_hash: &RHash,
+) -> Result<(String, bool, Option<String>), MagnusError> {
+    // extract `name:` if present, else default
+    let name_val = node_hash.get(Symbol::new("name"));
+    let name: String = if let Some(v) = name_val {
+        let name_option = parse_string_or_symbol(ruby, v)?;
+        name_option.unwrap_or_else(|| "".to_string())
+    } else {
+        "".to_string() // top-level might omit name
+    };
+
+    // extract `nullable:`
+    let nullable_val = node_hash.get(Symbol::new("nullable"));
+    let nullable: bool = if let Some(v) = nullable_val {
+        bool::try_convert(v).unwrap_or(true)
+    } else {
+        true // default to nullable
+    };
+
+    // optional `format:`
+    let format_val = node_hash.get(Symbol::new("format"));
+    let format: Option<String> = if let Some(v) = format_val {
+        parse_string_or_symbol(ruby, v)?
+    } else {
+        None
+    };
+
+    Ok((name, nullable, format))
+}
+
+/// Parse a struct schema node
+fn parse_struct_node(
+    ruby: &Ruby,
+    node_hash: &RHash,
+    name: String,
+    nullable: bool,
+) -> Result<SchemaNode, MagnusError> {
+    // parse subfields array from `fields`
+    let fields_val = node_hash.get(Symbol::new("fields")).ok_or_else(|| {
+        MagnusError::new(
+            ruby.exception_arg_error(),
+            "Struct must have :fields array defined",
+        )
+    })?;
+    let fields_arr: RArray = RArray::try_convert(fields_val).map_err(|_| {
+        MagnusError::new(
+            ruby.exception_type_error(),
+            "The :fields value must be an array",
+        )
+    })?;
+
+    // Check for empty struct immediately
+    if fields_arr.len() == 0 {
+        return Err(MagnusError::new(
+            ruby.exception_arg_error(),
+            format!("Cannot create a struct with zero fields. Struct name: '{}'. Parquet doesn't support empty structs", name)
+        ));
+    }
+
+    let mut fields = Vec::with_capacity(fields_arr.len());
+    for item in fields_arr.into_iter() {
+        fields.push(parse_schema_node(ruby, item)?);
+    }
+
+    Ok(SchemaNode::Struct {
+        name,
+        nullable,
+        fields,
+    })
+}
+
+/// Parse a list schema node
+fn parse_list_node(
+    ruby: &Ruby,
+    node_hash: &RHash,
+    name: String,
+    nullable: bool,
+) -> Result<SchemaNode, MagnusError> {
+    // parse `item`
+    let item_val = node_hash.get(Symbol::new("item")).ok_or_else(|| {
+        MagnusError::new(
+            ruby.exception_arg_error(),
+            "List type must have :item field defined",
+        )
+    })?;
+    let item_node = parse_schema_node(ruby, item_val)?;
+
+    Ok(SchemaNode::List {
+        name,
+        nullable,
+        item: Box::new(item_node),
+    })
+}
+
+/// Parse a map schema node
+fn parse_map_node(
+    ruby: &Ruby,
+    node_hash: &RHash,
+    name: String,
+    nullable: bool,
+) -> Result<SchemaNode, MagnusError> {
+    // parse `key` and `value`
+    let key_val = node_hash.get(Symbol::new("key")).ok_or_else(|| {
+        MagnusError::new(
+            ruby.exception_arg_error(),
+            "Map type must have :key field defined",
+        )
+    })?;
+    let value_val = node_hash.get(Symbol::new("value")).ok_or_else(|| {
+        MagnusError::new(
+            ruby.exception_arg_error(),
+            "Map type must have :value field defined",
+        )
+    })?;
+
+    let key_node = parse_schema_node(ruby, key_val)?;
+    let value_node = parse_schema_node(ruby, value_val)?;
+
+    Ok(SchemaNode::Map {
+        name,
+        nullable,
+        key: Box::new(key_node),
+        value: Box::new(value_node),
+    })
+}
+
+/// Parse a Ruby schema hash into a SchemaNode tree
+pub fn parse_schema_node(ruby: &Ruby, node_value: Value) -> Result<SchemaNode, MagnusError> {
+    // The node_value should be a Ruby Hash with keys: :name, :type, :nullable, etc.
+    let node_hash = RHash::from_value(node_value).ok_or_else(|| {
+        MagnusError::new(
+            ruby.exception_type_error(),
+            "Schema node must be a Hash with :type and other fields",
+        )
+    })?;
+
+    // extract `type:` which is a symbol/string
+    let type_val = node_hash.get(Symbol::new("type")).ok_or_else(|| {
+        MagnusError::new(
+            ruby.exception_arg_error(),
+            "Missing required :type field in schema node",
+        )
+    })?;
+    let type_str_option = parse_string_or_symbol(ruby, type_val)?;
+    let type_str = type_str_option.ok_or_else(|| {
+        MagnusError::new(
+            ruby.exception_arg_error(),
+            "Type cannot be nil - please specify a valid type string or symbol",
+        )
+    })?;
+
+    // Extract common fields (name, nullable, format)
+    let (name, nullable, format) = extract_common_fields(ruby, &node_hash)?;
+
+    // Delegate to type-specific parsers with clear error handling
+    match type_str.as_str() {
+        "struct" => parse_struct_node(ruby, &node_hash, name, nullable),
+        "list" => parse_list_node(ruby, &node_hash, name, nullable),
+        "map" => parse_map_node(ruby, &node_hash, name, nullable),
+        // For primitives, provide better error messages when type isn't recognized
+        other => {
+            if let Some(parquet_type) = parse_primitive_type(other) {
+                Ok(SchemaNode::Primitive {
+                    name,
+                    parquet_type,
+                    nullable,
+                    format,
+                })
+            } else {
+                Err(MagnusError::new(
+                    magnus::exception::arg_error(),
+                    format!(
+                        "Unknown type: '{}'. Supported types are: struct, list, map, int8, int16, int32, int64, uint8, uint16, uint32, uint64, float32, float64, boolean, string, binary, date32, timestamp_millis, timestamp_micros",
+                        other
+                    )
+                ))
+            }
+        }
+    }
+}
+
+/// Convert a type string like "int32" to a PrimitiveType
+fn parse_primitive_type(s: &str) -> Option<PrimitiveType> {
+    match s.to_lowercase().as_str() {
+        "int8" | "i8" => Some(PrimitiveType::Int8),
+        "int16" | "i16" => Some(PrimitiveType::Int16),
+        "int32" | "i32" | "int" => Some(PrimitiveType::Int32),
+        "int64" | "i64" | "long" | "bigint" => Some(PrimitiveType::Int64),
+        "uint8" | "u8" | "byte" => Some(PrimitiveType::UInt8),
+        "uint16" | "u16" => Some(PrimitiveType::UInt16),
+        "uint32" | "u32" | "uint" => Some(PrimitiveType::UInt32),
+        "uint64" | "u64" | "ulong" => Some(PrimitiveType::UInt64),
+        "float" | "float32" | "f32" => Some(PrimitiveType::Float32),
+        "double" | "float64" | "f64" => Some(PrimitiveType::Float64),
+        "bool" | "boolean" => Some(PrimitiveType::Boolean),
+        "string" | "utf8" | "str" | "text" => Some(PrimitiveType::String),
+        "binary" | "bytes" | "blob" => Some(PrimitiveType::Binary),
+        "date" | "date32" => Some(PrimitiveType::Date32),
+        "timestamp_millis" | "timestamp_ms" => Some(PrimitiveType::TimestampMillis),
+        "timestamp_micros" | "timestamp_us" => Some(PrimitiveType::TimestampMicros),
+        _ => None,
+    }
+}
+
+/// Convert a SchemaNode to an Arrow field
+pub fn schema_node_to_arrow_field(node: &SchemaNode) -> ArrowField {
+    match node {
+        SchemaNode::Primitive {
+            name,
+            parquet_type,
+            nullable,
+            format: _,
+        } => {
+            let dt = match parquet_type {
+                PrimitiveType::Int8 => ArrowDataType::Int8,
+                PrimitiveType::Int16 => ArrowDataType::Int16,
+                PrimitiveType::Int32 => ArrowDataType::Int32,
+                PrimitiveType::Int64 => ArrowDataType::Int64,
+                PrimitiveType::UInt8 => ArrowDataType::UInt8,
+                PrimitiveType::UInt16 => ArrowDataType::UInt16,
+                PrimitiveType::UInt32 => ArrowDataType::UInt32,
+                PrimitiveType::UInt64 => ArrowDataType::UInt64,
+                PrimitiveType::Float32 => ArrowDataType::Float32,
+                PrimitiveType::Float64 => ArrowDataType::Float64,
+                PrimitiveType::Boolean => ArrowDataType::Boolean,
+                PrimitiveType::String => ArrowDataType::Utf8,
+                PrimitiveType::Binary => ArrowDataType::Binary,
+                PrimitiveType::Date32 => ArrowDataType::Date32,
+                PrimitiveType::TimestampMillis => {
+                    ArrowDataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None)
+                }
+                PrimitiveType::TimestampMicros => {
+                    ArrowDataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None)
+                }
+            };
+            ArrowField::new(name, dt, *nullable)
+        }
+
+        SchemaNode::List {
+            name,
+            nullable,
+            item,
+        } => {
+            let child_field = schema_node_to_arrow_field(item);
+            let list_type = ArrowDataType::List(Arc::new(child_field));
+            ArrowField::new(name, list_type, *nullable)
+        }
+
+        SchemaNode::Map {
+            name,
+            nullable,
+            key,
+            value,
+        } => {
+            // A Map is basically: Map( Struct([key, value]), keysSorted=false )
+            let key_field = schema_node_to_arrow_field(key);
+            let value_field = schema_node_to_arrow_field(value);
+
+            let entries_struct = ArrowDataType::Struct(ArrowFields::from(vec![
+                ArrowField::new("key", key_field.data_type().clone(), false),
+                ArrowField::new(
+                    "value",
+                    value_field.data_type().clone(),
+                    value_field.is_nullable(),
+                ),
+            ]));
+
+            let map_data_type = ArrowDataType::Map(
+                Arc::new(ArrowField::new("entries", entries_struct, false)),
+                false, // not sorted
+            );
+            ArrowField::new(name, map_data_type, *nullable)
+        }
+
+        SchemaNode::Struct {
+            name,
+            nullable,
+            fields,
+        } => {
+            // Field validation happens earlier - no empty structs allowed
+            let mut arrow_subfields = Vec::with_capacity(fields.len());
+            for f in fields {
+                arrow_subfields.push(schema_node_to_arrow_field(f));
+            }
+            let struct_type = ArrowDataType::Struct(ArrowFields::from(arrow_subfields));
+            ArrowField::new(name, struct_type, *nullable)
+        }
+    }
+}
+
+/// Build an Arrow schema from the top-level Node, which must be a Struct
+pub fn build_arrow_schema(
+    root: &SchemaNode,
+    logger: &RubyLogger,
+) -> Result<Arc<ArrowSchema>, MagnusError> {
+    match root {
+        SchemaNode::Struct { fields, .. } => {
+            // Fields debug output removed - we've fixed the empty struct issue
+
+            let arrow_fields: Vec<ArrowField> =
+                fields.iter().map(schema_node_to_arrow_field).collect();
+            let arrow_schema = ArrowSchema::new(arrow_fields);
+            logger.debug(|| format!("Constructed Arrow schema: {:?}", arrow_schema))?;
+            Ok(Arc::new(arrow_schema))
+        }
+        _ => Err(MagnusError::new(
+            magnus::exception::arg_error(),
+            "Top-level schema must be a Struct".to_owned(),
+        )),
+    }
+}