parquet 0.3.0 → 0.3.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5416216b30f385f5ff8a9ee44a02eb1f685ac200c1bec340ad09e936a37ef9b6
-  data.tar.gz: ed9dd4781af2c5cab68d94e582cfd29135d37d399d7904b82dcffe4bfefbe52f
+  metadata.gz: f8fa4ebaf63f622e06e21085899249ea50edbceed3afa0c378df69bd0bdc8c58
+  data.tar.gz: 9f02946885241758cb26eff3564a6f386c4fafaa77811627e3af2f5768ea002a
 SHA512:
-  metadata.gz: 4f3650bac3f11755e691b7cf50f03cce66036a7e8989297c08aa6f0ab8cea60b3bcf0f6881e9a5e452ca51d78f64f183e9c20c46c38dad312bb568042ab7f0d4
-  data.tar.gz: 8a070d861e155d2d9b3cf31ec79d77307c379b03a7c797dd103a0e7837ff623a6507723621ef45715d9c8e764c836e786ab9a8a6ca3c5454551b3f1c3489018d
+  metadata.gz: 67edb7775baaec1f214edd359d1aff1768bceb1b24e3d616b23c2a28f60ae3a7c7d9c68b0a767d99558aa3841b725030452cb9726e0567c7e413f15c3da6562d
+  data.tar.gz: a704b96748fb53a272964fddcbe8db5ee4615fc4e96cec51e3897ab2e651b84c5429ebc2a50a48568da71e642b728c8f24b0aa46eab62f896f98a8919ec9141b

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

@@ -26,8 +26,9 @@ use crate::{
     IoLikeValue, ParquetSchemaType, ParquetWriteArgs, SchemaField, SendableWrite,
 };
 
+const MIN_SAMPLES_FOR_ESTIMATE: usize = 10; // Minimum samples needed for estimation
 const SAMPLE_SIZE: usize = 100; // Number of rows to sample for size estimation
-const MIN_BATCH_SIZE: usize = 100; // Minimum batch size to maintain efficiency
+const MIN_BATCH_SIZE: usize = 10; // Minimum batch size to maintain efficiency
 const INITIAL_BATCH_SIZE: usize = 100; // Initial batch size while sampling
 
 // Maximum memory usage per batch (64MB by default)
@@ -178,12 +179,6 @@ fn estimate_single_row_size(row: &RArray, schema: &[SchemaField]) -> Result<usiz
     Ok(row_size)
 }
 
-/// Calculate optimal batch size based on memory threshold and estimated row size
-fn calculate_batch_size(row_size: usize, memory_threshold: usize) -> usize {
-    let batch_size = memory_threshold / row_size;
-    batch_size.max(MIN_BATCH_SIZE)
-}
-
 #[inline]
 pub fn write_rows(args: &[Value]) -> Result<(), MagnusError> {
     let ruby = unsafe { Ruby::get_unchecked() };
@@ -293,10 +288,15 @@ pub fn write_rows(args: &[Value]) -> Result<(), MagnusError> {
                     rows_in_batch += 1;
                     total_rows += 1;
 
-                    // Recalculate batch size if we have enough samples and no user-specified size
-                    if size_samples.len() >= sample_size && user_batch_size.is_none() {
-                        let avg_row_size = size_samples.iter().sum::<usize>() / size_samples.len();
-                        current_batch_size = calculate_batch_size(avg_row_size, flush_threshold);
+                    // Calculate batch size progressively once we have minimum samples
+                    if size_samples.len() >= MIN_SAMPLES_FOR_ESTIMATE && user_batch_size.is_none() {
+                        let total_size = size_samples.iter().sum::<usize>();
+                        // Safe because we know we have at least MIN_SAMPLES_FOR_ESTIMATE samples
+                        let avg_row_size = total_size as f64 / size_samples.len() as f64;
+                        let avg_row_size = avg_row_size.max(1.0); // Ensure we don't divide by zero
+                        let suggested_batch_size =
+                            (flush_threshold as f64 / avg_row_size).floor() as usize;
+                        current_batch_size = suggested_batch_size.max(MIN_BATCH_SIZE);
                     }
 
                     // When we reach batch size, write the batch

data/lib/parquet/version.rb

@@ -1,3 +1,3 @@
 module Parquet
-  VERSION = "0.3.0"
+  VERSION = "0.3.1"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: parquet
 version: !ruby/object:Gem::Version
-  version: 0.3.0
+  version: 0.3.1
 platform: ruby
 authors:
 - Nathan Jaremko