parquet 0.5.12 → 0.6.0

data/ext/parquet-core/tests/writer_tests.rs

@@ -0,0 +1,545 @@
+use bytes::Bytes;
+use ordered_float::OrderedFloat;
+use parquet::basic::Compression;
+use parquet::file::properties::WriterProperties;
+use parquet_core::*;
+use std::sync::Arc;
+
+mod test_helpers;
+use test_helpers::*;
+
+// =============================================================================
+// Basic Writer Functionality Tests
+// =============================================================================
+
+#[test]
+fn test_writer_basic_functionality() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![
+                SchemaNode::Primitive {
+                    name: "id".to_string(),
+                    primitive_type: PrimitiveType::Int32,
+                    nullable: false,
+                    format: None,
+                },
+                SchemaNode::Primitive {
+                    name: "name".to_string(),
+                    primitive_type: PrimitiveType::String,
+                    nullable: true,
+                    format: None,
+                },
+            ],
+        })
+        .build()
+        .unwrap();
+
+    let rows = vec![
+        vec![ParquetValue::Int32(1), ParquetValue::String("Alice".into())],
+        vec![
+            ParquetValue::Int32(2),
+            ParquetValue::Null, // nullable field
+        ],
+    ];
+
+    test_roundtrip(rows, schema).unwrap();
+}
+
+// =============================================================================
+// Batch Size Configuration Tests
+// =============================================================================
+
+#[test]
+fn test_writer_fixed_batch_sizes() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![
+                SchemaNode::Primitive {
+                    name: "id".to_string(),
+                    primitive_type: PrimitiveType::Int64,
+                    nullable: false,
+                    format: None,
+                },
+                SchemaNode::Primitive {
+                    name: "data".to_string(),
+                    primitive_type: PrimitiveType::String,
+                    nullable: false,
+                    format: None,
+                },
+            ],
+        })
+        .build()
+        .unwrap();
+
+    // Test different batch sizes
+    let batch_sizes = vec![10, 100, 1000, 5000];
+
+    for batch_size in batch_sizes {
+        // Generate test data
+        let rows: Vec<Vec<ParquetValue>> = (0..10000)
+            .map(|i| {
+                vec![
+                    ParquetValue::Int64(i),
+                    ParquetValue::String(Arc::from(format!("Row {}", i))),
+                ]
+            })
+            .collect();
+
+        // Use test_roundtrip_with_options for batch size testing
+        let result = test_roundtrip_with_options(
+            rows,
+            schema.clone(),
+            Compression::UNCOMPRESSED,
+            Some(batch_size),
+        );
+
+        assert!(
+            result.is_ok(),
+            "Batch size {} failed: {:?}",
+            batch_size,
+            result
+        );
+    }
+}
+
+#[test]
+fn test_writer_adaptive_batch_sizing() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![
+                SchemaNode::Primitive {
+                    name: "id".to_string(),
+                    primitive_type: PrimitiveType::Int32,
+                    nullable: false,
+                    format: None,
+                },
+                SchemaNode::Primitive {
+                    name: "variable_string".to_string(),
+                    primitive_type: PrimitiveType::String,
+                    nullable: false,
+                    format: None,
+                },
+            ],
+        })
+        .build()
+        .unwrap();
+
+    let mut buffer = Vec::new();
+    {
+        // Don't set a fixed batch size - let it adapt
+        let mut writer = WriterBuilder::new()
+            .with_sample_size(50)
+            .build(&mut buffer, schema)
+            .unwrap();
+
+        // Write rows with varying sizes
+        for i in 0..1000 {
+            let string_size = if i % 100 == 0 {
+                10000 // Large string every 100 rows
+            } else {
+                100 // Normal string
+            };
+
+            let row = vec![
+                ParquetValue::Int32(i),
+                ParquetValue::String(Arc::from("x".repeat(string_size))),
+            ];
+
+            writer.write_row(row).unwrap();
+        }
+
+        writer.close().unwrap();
+    }
+
+    // Verify all data was written
+    let bytes = Bytes::from(buffer);
+    let reader = Reader::new(bytes);
+
+    let read_rows: Vec<_> = reader
+        .read_rows()
+        .unwrap()
+        .collect::<Result<Vec<_>>>()
+        .unwrap();
+
+    assert_eq!(read_rows.len(), 1000);
+
+    // Verify variable string sizes
+    for (i, row) in read_rows.iter().enumerate() {
+        match &row[1] {
+            ParquetValue::String(s) => {
+                let expected_len = if i % 100 == 0 { 10000 } else { 100 };
+                assert_eq!(s.len(), expected_len, "Wrong string length at row {}", i);
+            }
+            _ => panic!("Expected string value"),
+        }
+    }
+}
+
+// =============================================================================
+// Memory Management Tests
+// =============================================================================
+
+#[test]
+fn test_memory_threshold_configuration() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![SchemaNode::Primitive {
+                name: "large_string".to_string(),
+                primitive_type: PrimitiveType::String,
+                nullable: false,
+                format: None,
+            }],
+        })
+        .build()
+        .unwrap();
+
+    // Test with different memory thresholds
+    let thresholds = vec![
+        1024 * 1024,      // 1MB
+        10 * 1024 * 1024, // 10MB
+        50 * 1024 * 1024, // 50MB
+    ];
+
+    for threshold in thresholds {
+        let mut buffer = Vec::new();
+        {
+            let mut writer = WriterBuilder::new()
+                .with_memory_threshold(threshold)
+                .build(&mut buffer, schema.clone())
+                .unwrap();
+
+            // Write large strings that will trigger memory-based flushing
+            let large_string: Arc<str> = Arc::from("x".repeat(1024)); // 1KB string
+            let rows: Vec<Vec<ParquetValue>> = (0..5000)
+                .map(|_| vec![ParquetValue::String(large_string.clone())])
+                .collect();
+
+            writer.write_rows(rows).unwrap();
+            writer.close().unwrap();
+        }
+
+        // Verify data was written correctly
+        let bytes = Bytes::from(buffer);
+        let reader = Reader::new(bytes);
+        let read_count = reader.read_rows().unwrap().count();
+        assert_eq!(read_count, 5000);
+    }
+}
+
+#[test]
+fn test_writer_memory_flushing_with_binary() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![SchemaNode::Primitive {
+                name: "data".to_string(),
+                primitive_type: PrimitiveType::Binary,
+                nullable: true,
+                format: None,
+            }],
+        })
+        .build()
+        .unwrap();
+
+    // Generate test data
+    let rows: Vec<Vec<ParquetValue>> = (0..100)
+        .map(|i| {
+            let size = if i % 10 == 0 { 500 } else { 50 };
+            vec![ParquetValue::Bytes(Bytes::from(vec![i as u8; size]))]
+        })
+        .collect();
+
+    // Use a custom writer with memory threshold
+    let mut buffer = Vec::new();
+    {
+        let mut writer = WriterBuilder::new()
+            .with_memory_threshold(1024) // 1KB threshold
+            .with_sample_size(5)
+            .build(&mut buffer, schema.clone())
+            .unwrap();
+
+        writer.write_rows(rows.clone()).unwrap();
+        writer.close().unwrap();
+    }
+
+    // Verify data was written correctly
+    let reader = Reader::new(Bytes::from(buffer));
+    let read_rows: Vec<_> = reader
+        .read_rows()
+        .unwrap()
+        .collect::<Result<Vec<_>>>()
+        .unwrap();
+
+    assert_eq!(read_rows, rows);
+}
+
+// =============================================================================
+// Advanced Configuration Tests
+// =============================================================================
+
+#[test]
+fn test_writer_properties_direct() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![SchemaNode::Primitive {
+                name: "value".to_string(),
+                primitive_type: PrimitiveType::String,
+                nullable: false,
+                format: None,
+            }],
+        })
+        .build()
+        .unwrap();
+
+    // Test custom writer properties
+    let props = WriterProperties::builder()
+        .set_writer_version(parquet::file::properties::WriterVersion::PARQUET_2_0)
+        .set_compression(Compression::ZSTD(
+            parquet::basic::ZstdLevel::try_new(3).unwrap(),
+        ))
+        .set_data_page_size_limit(1024) // Small page size
+        .set_dictionary_enabled(true)
+        .set_statistics_enabled(parquet::file::properties::EnabledStatistics::Page)
+        .build();
+
+    let mut buffer = Vec::new();
+    {
+        let mut writer = Writer::new_with_properties(&mut buffer, schema, props).unwrap();
+
+        // Write data with repeated values to test dictionary encoding
+        let rows: Vec<Vec<ParquetValue>> = (0..1000)
+            .map(|i| {
+                vec![ParquetValue::String(Arc::from(format!(
+                    "Category_{}",
+                    i % 10
+                )))]
+            })
+            .collect();
+
+        writer.write_rows(rows).unwrap();
+        writer.close().unwrap();
+    }
+
+    // Read back and verify
+    let bytes = Bytes::from(buffer);
+    let mut reader = Reader::new(bytes);
+
+    // Check metadata
+    let metadata = reader.metadata().unwrap();
+    assert!(metadata.num_rows() == 1000);
+
+    // Verify data integrity
+    let read_count = reader.read_rows().unwrap().count();
+    assert_eq!(read_count, 1000);
+}
+
+#[test]
+fn test_writer_version_compatibility() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![SchemaNode::Primitive {
+                name: "value".to_string(),
+                primitive_type: PrimitiveType::Int32,
+                nullable: false,
+                format: None,
+            }],
+        })
+        .build()
+        .unwrap();
+
+    let rows: Vec<Vec<ParquetValue>> = (0..100).map(|i| vec![ParquetValue::Int32(i)]).collect();
+
+    // Test different writer versions
+    let versions = vec![
+        parquet::file::properties::WriterVersion::PARQUET_1_0,
+        parquet::file::properties::WriterVersion::PARQUET_2_0,
+    ];
+
+    for version in versions {
+        let mut buffer = Vec::new();
+        {
+            let props = WriterProperties::builder()
+                .set_writer_version(version)
+                .build();
+
+            let mut writer =
+                Writer::new_with_properties(&mut buffer, schema.clone(), props).unwrap();
+            writer.write_rows(rows.clone()).unwrap();
+            writer.close().unwrap();
+        }
+
+        // Verify we can read both versions
+        let bytes = Bytes::from(buffer);
+        let reader = Reader::new(bytes);
+
+        let read_rows: Vec<_> = reader
+            .read_rows()
+            .unwrap()
+            .collect::<Result<Vec<_>>>()
+            .unwrap();
+
+        assert_eq!(read_rows.len(), 100);
+    }
+}
+
+// =============================================================================
+// Large Data Handling Tests
+// =============================================================================
+
+#[test]
+fn test_large_string_handling() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![
+                SchemaNode::Primitive {
+                    name: "id".to_string(),
+                    primitive_type: PrimitiveType::Int32,
+                    nullable: false,
+                    format: None,
+                },
+                SchemaNode::Primitive {
+                    name: "content".to_string(),
+                    primitive_type: PrimitiveType::String,
+                    nullable: true,
+                    format: None,
+                },
+            ],
+        })
+        .build()
+        .unwrap();
+
+    // Create strings of various sizes
+    let small = "a".repeat(100);
+    let medium = "b".repeat(10_000);
+    let large = "c".repeat(100_000);
+
+    // Generate test data
+    let rows: Vec<Vec<ParquetValue>> = (0..30)
+        .map(|i| {
+            let content = match i % 3 {
+                0 => ParquetValue::String(small.clone().into()),
+                1 => ParquetValue::String(medium.clone().into()),
+                2 => ParquetValue::String(large.clone().into()),
+                _ => unreachable!(),
+            };
+
+            vec![ParquetValue::Int32(i), content]
+        })
+        .collect();
+
+    // Use custom writer with memory threshold
+    let mut buffer = Vec::new();
+    {
+        let mut writer = WriterBuilder::new()
+            .with_memory_threshold(1024 * 1024) // 1MB
+            .build(&mut buffer, schema.clone())
+            .unwrap();
+
+        writer.write_rows(rows.clone()).unwrap();
+        writer.close().unwrap();
+    }
+
+    // Verify all data was written
+    let reader = Reader::new(Bytes::from(buffer));
+    let read_rows: Vec<_> = reader.read_rows().unwrap().collect::<Result<_>>().unwrap();
+    assert_eq!(read_rows, rows);
+}
+
+#[test]
+fn test_complex_nested_data_memory() {
+    let schema = SchemaBuilder::new()
+        .with_root(SchemaNode::Struct {
+            name: "root".to_string(),
+            nullable: false,
+            fields: vec![
+                SchemaNode::Primitive {
+                    name: "id".to_string(),
+                    primitive_type: PrimitiveType::Int32,
+                    nullable: false,
+                    format: None,
+                },
+                SchemaNode::List {
+                    name: "items".to_string(),
+                    nullable: true,
+                    item: Box::new(SchemaNode::Struct {
+                        name: "item".to_string(),
+                        nullable: false,
+                        fields: vec![
+                            SchemaNode::Primitive {
+                                name: "key".to_string(),
+                                primitive_type: PrimitiveType::String,
+                                nullable: false,
+                                format: None,
+                            },
+                            SchemaNode::Primitive {
+                                name: "value".to_string(),
+                                primitive_type: PrimitiveType::Float64,
+                                nullable: true,
+                                format: None,
+                            },
+                        ],
+                    }),
+                },
+            ],
+        })
+        .build()
+        .unwrap();
+
+    // Generate test data
+    let rows: Vec<Vec<ParquetValue>> = (0..100)
+        .map(|i| {
+            let num_items = (i % 10 + 1) as usize;
+            let mut items = Vec::new();
+
+            for j in 0..num_items {
+                items.push(ParquetValue::Record(indexmap::indexmap! {
+                    "key".into() => ParquetValue::String(format!("key_{}_{}", i, j).into()),
+                    "value".into() => if j % 2 == 0 {
+                        ParquetValue::Float64(OrderedFloat(j as f64 * 1.5))
+                    } else {
+                        ParquetValue::Null
+                    },
+                }));
+            }
+
+            vec![ParquetValue::Int32(i), ParquetValue::List(items)]
+        })
+        .collect();
+
+    // Use custom writer with memory threshold
+    let mut buffer = Vec::new();
+    {
+        let mut writer = WriterBuilder::new()
+            .with_memory_threshold(500 * 1024) // 500KB
+            .build(&mut buffer, schema.clone())
+            .unwrap();
+
+        writer.write_rows(rows.clone()).unwrap();
+        writer.close().unwrap();
+    }
+
+    // Read back and verify
+    let reader = Reader::new(Bytes::from(buffer));
+    let read_rows: Vec<_> = reader
+        .read_rows()
+        .unwrap()
+        .collect::<Result<Vec<_>>>()
+        .unwrap();
+
+    assert_eq!(read_rows, rows);
+}

data/ext/parquet-ruby-adapter/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "parquet-ruby-adapter"
+version = "0.1.0"
+edition = "2021"
+
+[build-dependencies]
+rb-sys-env = "^0.2"
+
+[dependencies]
+arrow-array = { git = "https://github.com/njaremko/arrow-rs", branch = "nathan_06-24-remove_primitive_map_key_assertion_on_record_reader" }
+arrow-buffer = { git = "https://github.com/njaremko/arrow-rs", branch = "nathan_06-24-remove_primitive_map_key_assertion_on_record_reader" }
+arrow-schema = { git = "https://github.com/njaremko/arrow-rs", branch = "nathan_06-24-remove_primitive_map_key_assertion_on_record_reader" }
+bytes = "1.5"
+magnus = { version = "0.7", features = ["rb-sys"] }
+num = "0.4.3"
+ordered-float = "5.0.0"
+parquet = { git = "https://github.com/njaremko/arrow-rs", branch = "nathan_06-24-remove_primitive_map_key_assertion_on_record_reader", features = ["arrow"] }
+parquet-core = { path = "../parquet-core" }
+rb-sys = { version = "0.9", features = ["stable-api-compiled-fallback"] }
+tempfile = "^3.15"
+thiserror = "2.0"
+indexmap = "2.2"

data/ext/parquet-ruby-adapter/build.rs

@@ -0,0 +1,5 @@
+pub fn main() -> Result<(), Box<dyn std::error::Error>> {
+    rb_sys_env::activate()?;
+
+    Ok(())
+}

data/ext/parquet-ruby-adapter/examples/try_into_value_demo.rs

@@ -0,0 +1,98 @@
+use magnus::{IntoValue, Ruby, Value};
+use parquet_ruby_adapter::{Result, RubyAdapterError, TryIntoValue};
+
+/// Example struct that can fail during conversion to Ruby
+struct ComplexData {
+    name: String,
+    values: Vec<i32>,
+    metadata: std::collections::HashMap<String, String>,
+}
+
+impl TryIntoValue for ComplexData {
+    fn try_into_value(self, handle: &Ruby) -> Result<Value> {
+        let hash = handle.hash_new();
+
+        // Set name
+        hash.aset("name", self.name)
+            .map_err(|e| RubyAdapterError::type_conversion(format!("Failed to set name: {}", e)))?;
+
+        // Convert values array
+        let values_array = handle.ary_new();
+        for value in self.values {
+            values_array.push(value).map_err(|e| {
+                RubyAdapterError::type_conversion(format!("Failed to push value: {}", e))
+            })?;
+        }
+        hash.aset("values", values_array).map_err(|e| {
+            RubyAdapterError::type_conversion(format!("Failed to set values: {}", e))
+        })?;
+
+        // Convert metadata hash
+        let metadata_hash = handle.hash_new();
+        for (key, value) in self.metadata {
+            metadata_hash.aset(key.clone(), value).map_err(|e| {
+                RubyAdapterError::type_conversion(format!(
+                    "Failed to set metadata key {}: {}",
+                    key, e
+                ))
+            })?;
+        }
+        hash.aset("metadata", metadata_hash).map_err(|e| {
+            RubyAdapterError::type_conversion(format!("Failed to set metadata: {}", e))
+        })?;
+
+        Ok(handle.into_value(hash))
+    }
+}
+
+// Example of a type that might fail validation during conversion
+struct ValidatedNumber {
+    value: i32,
+}
+
+impl TryIntoValue for ValidatedNumber {
+    fn try_into_value(self, handle: &Ruby) -> Result<Value> {
+        // Validate the number is positive
+        if self.value < 0 {
+            return Err(RubyAdapterError::type_conversion(format!(
+                "ValidatedNumber must be positive, got {}",
+                self.value
+            )));
+        }
+
+        // If valid, convert to Ruby
+        Ok(self.value.into_value_with(handle))
+    }
+}
+
+fn main() -> Result<()> {
+    // Example usage:
+    let ruby = Ruby::get().map_err(|_| RubyAdapterError::runtime("Failed to get Ruby runtime"))?;
+
+    // Success case
+    let data = ComplexData {
+        name: "example".to_string(),
+        values: vec![1, 2, 3],
+        metadata: std::collections::HashMap::from([
+            ("key1".to_string(), "value1".to_string()),
+            ("key2".to_string(), "value2".to_string()),
+        ]),
+    };
+
+    let _ruby_value = data.try_into_value(&ruby)?;
+    println!("Successfully converted ComplexData to Ruby value");
+
+    // Validation failure case
+    let invalid_number = ValidatedNumber { value: -5 };
+    match invalid_number.try_into_value(&ruby) {
+        Ok(_) => println!("This shouldn't happen"),
+        Err(e) => println!("Expected validation error: {}", e),
+    }
+
+    // Using the convenience method
+    let valid_number = ValidatedNumber { value: 42 };
+    let _ruby_value = valid_number.try_into_value_with_current_thread()?;
+    println!("Successfully converted ValidatedNumber to Ruby value");
+
+    Ok(())
+}