parquet 0.0.4 → 0.2.5

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

@@ -0,0 +1,204 @@
+use itertools::Itertools;
+
+use super::*;
+
+#[derive(Debug)]
+pub enum RowRecord<S: BuildHasher + Default> {
+    Vec(Vec<ParquetField>),
+    Map(HashMap<StringCacheKey, ParquetField, S>),
+}
+
+#[derive(Debug)]
+pub enum ColumnRecord<S: BuildHasher + Default> {
+    Vec(Vec<Vec<ParquetValue>>),
+    Map(HashMap<StringCacheKey, Vec<ParquetValue>, S>),
+}
+
+#[derive(Debug)]
+pub struct ParquetField(pub Field);
+
+impl<S: BuildHasher + Default> IntoValue for RowRecord<S> {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        match self {
+            RowRecord::Vec(vec) => {
+                let ary = handle.ary_new_capa(vec.len());
+                vec.into_iter().try_for_each(|v| ary.push(v)).unwrap();
+                handle.into_value(ary)
+            }
+            RowRecord::Map(map) => {
+                let hash = handle.hash_new_capa(map.len());
+
+                let mut values: [Value; 128] = [handle.qnil().as_value(); 128];
+                let mut i = 0;
+
+                for chunk in &map.into_iter().chunks(64) {
+                    // Reduced to 64 to ensure space for pairs
+                    for (k, v) in chunk {
+                        if i + 1 >= values.len() {
+                            // Bulk insert current batch if array is full
+                            hash.bulk_insert(&values[..i]).unwrap();
+                            values[..i].fill(handle.qnil().as_value());
+                            i = 0;
+                        }
+                        values[i] = handle.into_value(k);
+                        values[i + 1] = handle.into_value(v);
+                        i += 2;
+                    }
+                    // Insert any remaining pairs
+                    if i > 0 {
+                        hash.bulk_insert(&values[..i]).unwrap();
+                        values[..i].fill(handle.qnil().as_value());
+                        i = 0;
+                    }
+                }
+
+                hash.into_value_with(handle)
+            }
+        }
+    }
+}
+
+impl<S: BuildHasher + Default> IntoValue for ColumnRecord<S> {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        match self {
+            ColumnRecord::Vec(vec) => {
+                let ary = handle.ary_new_capa(vec.len());
+                vec.into_iter()
+                    .try_for_each(|v| {
+                        let nested_ary = handle.ary_new_capa(v.len());
+                        v.into_iter().try_for_each(|v| nested_ary.push(v)).unwrap();
+                        ary.push(nested_ary.into_value_with(handle))
+                    })
+                    .unwrap();
+                ary.into_value_with(handle)
+            }
+            ColumnRecord::Map(map) => {
+                let hash = handle.hash_new_capa(map.len());
+
+                let mut values: [Value; 128] = [handle.qnil().as_value(); 128];
+                let mut i = 0;
+
+                for chunk in &map.into_iter().chunks(64) {
+                    // Reduced to 64 to ensure space for pairs
+                    for (k, v) in chunk {
+                        if i + 1 >= values.len() {
+                            // Bulk insert current batch if array is full
+                            hash.bulk_insert(&values[..i]).unwrap();
+                            values[..i].fill(handle.qnil().as_value());
+                            i = 0;
+                        }
+                        values[i] = handle.into_value(k);
+                        let ary = handle.ary_new_capa(v.len());
+                        v.into_iter().try_for_each(|v| ary.push(v)).unwrap();
+                        values[i + 1] = handle.into_value(ary);
+                        i += 2;
+                    }
+                    // Insert any remaining pairs
+                    if i > 0 {
+                        hash.bulk_insert(&values[..i]).unwrap();
+                        values[..i].fill(handle.qnil().as_value());
+                        i = 0;
+                    }
+                }
+
+                hash.into_value_with(handle)
+            }
+        }
+    }
+}
+
+impl IntoValue for ParquetField {
+    fn into_value_with(self, handle: &Ruby) -> Value {
+        match self.0 {
+            Field::Null => handle.qnil().as_value(),
+            Field::Bool(b) => b.into_value_with(handle),
+            Field::Short(s) => s.into_value_with(handle),
+            Field::Int(i) => i.into_value_with(handle),
+            Field::Long(l) => l.into_value_with(handle),
+            Field::UByte(ub) => ub.into_value_with(handle),
+            Field::UShort(us) => us.into_value_with(handle),
+            Field::UInt(ui) => ui.into_value_with(handle),
+            Field::ULong(ul) => ul.into_value_with(handle),
+            Field::Float16(f) => f32::from(f).into_value_with(handle),
+            Field::Float(f) => f.into_value_with(handle),
+            Field::Double(d) => d.into_value_with(handle),
+            Field::Str(s) => s.into_value_with(handle),
+            Field::Byte(b) => b.into_value_with(handle),
+            Field::Bytes(b) => handle.str_from_slice(b.data()).as_value(),
+            Field::Date(d) => {
+                let ts = jiff::Timestamp::from_second((d as i64) * 86400).unwrap();
+                let formatted = ts.strftime("%Y-%m-%d").to_string();
+                formatted.into_value_with(handle)
+            }
+            Field::TimestampMillis(ts) => {
+                let ts = jiff::Timestamp::from_millisecond(ts).unwrap();
+                let time_class = handle.class_time();
+                time_class
+                    .funcall::<_, _, Value>("parse", (ts.to_string(),))
+                    .unwrap()
+                    .into_value_with(handle)
+            }
+            Field::TimestampMicros(ts) => {
+                let ts = jiff::Timestamp::from_microsecond(ts).unwrap();
+                let time_class = handle.class_time();
+                time_class
+                    .funcall::<_, _, Value>("parse", (ts.to_string(),))
+                    .unwrap()
+                    .into_value_with(handle)
+            }
+            Field::ListInternal(list) => {
+                let elements = list.elements();
+                let ary = handle.ary_new_capa(elements.len());
+                elements
+                    .iter()
+                    .try_for_each(|e| ary.push(ParquetField(e.clone()).into_value_with(handle)))
+                    .unwrap();
+                ary.into_value_with(handle)
+            }
+            Field::MapInternal(map) => {
+                let entries = map.entries();
+                let hash = handle.hash_new_capa(entries.len());
+                entries
+                    .iter()
+                    .try_for_each(|(k, v)| {
+                        hash.aset(
+                            ParquetField(k.clone()).into_value_with(handle),
+                            ParquetField(v.clone()).into_value_with(handle),
+                        )
+                    })
+                    .unwrap();
+                hash.into_value_with(handle)
+            }
+            Field::Decimal(d) => {
+                let value = match d {
+                    Decimal::Int32 { value, scale, .. } => {
+                        let unscaled = i32::from_be_bytes(value);
+                        format!("{}e-{}", unscaled, scale)
+                    }
+                    Decimal::Int64 { value, scale, .. } => {
+                        let unscaled = i64::from_be_bytes(value);
+                        format!("{}e-{}", unscaled, scale)
+                    }
+                    Decimal::Bytes { value, scale, .. } => {
+                        // Convert bytes to string representation of unscaled value
+                        let unscaled = String::from_utf8_lossy(value.data());
+                        format!("{}e-{}", unscaled, scale)
+                    }
+                };
+                handle.eval(&format!("BigDecimal(\"{value}\")")).unwrap()
+            }
+            Field::Group(row) => {
+                let hash = handle.hash_new();
+                row.get_column_iter()
+                    .try_for_each(|(k, v)| {
+                        hash.aset(
+                            k.clone().into_value_with(handle),
+                            ParquetField(v.clone()).into_value_with(handle),
+                        )
+                    })
+                    .unwrap();
+                hash.into_value_with(handle)
+            }
+        }
+    }
+}

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

@@ -0,0 +1,85 @@
+use super::*;
+
+pub fn parse_zoned_timestamp(value: &ParquetValue) -> jiff::Timestamp {
+    let (ts, tz) = match value {
+        ParquetValue::TimestampSecond(ts, tz) => (jiff::Timestamp::from_second(*ts).unwrap(), tz),
+        ParquetValue::TimestampMillis(ts, tz) => {
+            (jiff::Timestamp::from_millisecond(*ts).unwrap(), tz)
+        }
+        ParquetValue::TimestampMicros(ts, tz) => {
+            (jiff::Timestamp::from_microsecond(*ts).unwrap(), tz)
+        }
+        ParquetValue::TimestampNanos(ts, tz) => {
+            (jiff::Timestamp::from_nanosecond(*ts as i128).unwrap(), tz)
+        }
+        _ => panic!("Invalid timestamp value"),
+    };
+
+    // If timezone is provided, convert to zoned timestamp
+    if let Some(tz) = tz {
+        // Handle fixed offset timezones like "+09:00" first
+        if tz.starts_with('+') || tz.starts_with('-') {
+            // Parse the offset string into hours and minutes
+            let (hours, minutes) = if tz.len() >= 5 && tz.contains(':') {
+                // Format: "+09:00" or "-09:00"
+                let h = tz[1..3].parse::<i32>().unwrap_or(0);
+                let m = tz[4..6].parse::<i32>().unwrap_or(0);
+                (h, m)
+            } else if tz.len() >= 3 {
+                // Format: "+09" or "-09"
+                let h = tz[1..3].parse::<i32>().unwrap_or(0);
+                (h, 0)
+            } else {
+                (0, 0)
+            };
+
+            // Apply sign
+            let total_minutes = if tz.starts_with('-') {
+                -(hours * 60 + minutes)
+            } else {
+                hours * 60 + minutes
+            };
+
+            // Create fixed timezone
+            let tz = jiff::tz::TimeZone::fixed(jiff::tz::offset((total_minutes / 60) as i8));
+            ts.to_zoned(tz).timestamp()
+        } else {
+            // Try IANA timezone
+            match ts.intz(&tz) {
+                Ok(zoned) => zoned.timestamp(),
+                Err(_) => ts, // Fall back to UTC if timezone is invalid
+            }
+        }
+    } else {
+        // No timezone provided - treat as UTC
+        ts
+    }
+}
+
+// Macro for handling timestamp conversions
+#[macro_export]
+macro_rules! impl_timestamp_conversion {
+    ($value:expr, $unit:ident, $handle:expr) => {{
+        match $value {
+            ParquetValue::$unit(ts, tz) => {
+                let ts = parse_zoned_timestamp(&ParquetValue::$unit(ts, tz));
+                let time_class = $handle.class_time();
+                time_class
+                    .funcall::<_, _, Value>("parse", (ts.to_string(),))
+                    .unwrap()
+                    .into_value_with($handle)
+            }
+            _ => panic!("Invalid timestamp type"),
+        }
+    }};
+}
+
+// Macro for handling date conversions
+#[macro_export]
+macro_rules! impl_date_conversion {
+    ($value:expr, $handle:expr) => {{
+        let ts = jiff::Timestamp::from_second(($value as i64) * 86400).unwrap();
+        let formatted = ts.strftime("%Y-%m-%d").to_string();
+        formatted.into_value_with($handle)
+    }};
+}