From f18699995e683bccaf16caff4d92e3f56e740d02 Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Thu, 10 Sep 2026 01:58:31 +0100 Subject: [PATCH 1/7] Migrate Decimal Scalar for Minarrow 0.18.1 --- python/Cargo.lock | 17 ++++++++--------- python/Cargo.toml | 6 +++--- python/pyproject.toml | 4 ++-- rust/Cargo.toml | 5 +++-- rust/src/models/readers/ipc/window.rs | 27 +++++++++++++++++++++++++++ rust/src/models/types/parquet.rs | 6 ++++++ 6 files changed, 49 insertions(+), 16 deletions(-) diff --git a/python/Cargo.lock b/python/Cargo.lock index f394dbb..dc13912 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -662,7 +662,7 @@ checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "lightstream" -version = "0.6.0" +version = "0.6.1" dependencies = [ "bytes", "fast-float2", @@ -692,7 +692,7 @@ dependencies = [ [[package]] name = "lightstream-py" -version = "0.6.0" +version = "0.6.1" dependencies = [ "futures-core", "lightstream", @@ -734,11 +734,10 @@ checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "minarrow" -version = "0.16.2" +version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec048e9b7a43abce573068b7f4e220b82a78e44d19c2888e31527efa7a4be109" +checksum = "e53f5031b7d16aa76bc8671171814419694c687f5160917a3668f010dab2b609" dependencies = [ - "libc", "log", "num-traits", "vec64", @@ -746,9 +745,9 @@ dependencies = [ [[package]] name = "minarrow-pyo3" -version = "0.16.0" +version = "0.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cde6692dbbba1430e1205cb65c56420d3ff2440592ce982d98f2794a14a7b8ef" +checksum = "b6363f738664b325aab048cee8101e79606a487d9a8a294b960a1f55220961c4" dependencies = [ "minarrow", "pyo3", @@ -1674,9 +1673,9 @@ dependencies = [ [[package]] name = "vec64" -version = "0.4.7" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e14d6d15c0dbb3d2929a916fe20a15e7d04e787c9fbc29b3940e0654b11482c6" +checksum = "1aef6bbef159f21ac387220ebc71141524423f7886afd390bc7b140f12630425" [[package]] name = "version_check" diff --git a/python/Cargo.toml b/python/Cargo.toml index bf4ff3e..1463769 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -2,7 +2,7 @@ cargo-features = ["trim-paths"] [package] name = "lightstream-py" -version = "0.6.0" +version = "0.6.1" edition = "2024" authors = ["Peter G. Bower"] license = "MPL-2.0" @@ -24,8 +24,8 @@ lightstream = { version = "0.6", path = "../rust", features = ["csv", "datetime" # minarrow-pyo3's dictionary-index conversion needs the extended features, # and they flow through lightstream's flags so its match arms gate in step # with minarrow's variants. -minarrow = { version = "0.16", features = ["chunked"] } -minarrow-pyo3 = { version = "0.16", features = ["extended_categorical", "extended_numeric_types"] } +minarrow = { version = "0.18.1", features = ["chunked"] } +minarrow-pyo3 = { version = "0.18.1", features = ["extended_categorical", "extended_numeric_types"] } futures-core = "0.3" pyo3 = { version = "0.29", features = ["abi3-py39"] } # The QUIC and WebTransport dependency pins mirror lightstream's, so diff --git a/python/pyproject.toml b/python/pyproject.toml index c2ae316..400e7e8 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "lightstream-io" -version = "0.6.0" +version = "0.6.1" description = "Streaming Arrow I/O for Python - files, sockets, and network transports with zero-copy minarrow interop." readme = "README.md" requires-python = ">=3.9" @@ -20,7 +20,7 @@ classifiers = [ "Topic :: Scientific/Engineering", "Topic :: Software Development :: Libraries", ] -dependencies = ["minarrow>=0.16"] +dependencies = ["minarrow>=0.18"] [project.urls] Homepage = "https://github.com/SpaceCell/lightstream" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 04e07f4..e7a62a3 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -42,8 +42,8 @@ flatbuffers = "25.12.19" libc = { version = "0.2.183", optional = true } futures-sink = "0.3.32" log = "0.4" -minarrow = { version = "0.17.0", features = ["chunked", "views", "select", "size"], default-features = false } -vec64 = { version = "0.5.0" } +minarrow = { version = "0.18.1", features = ["chunked", "views", "select", "size"], default-features = false } +vec64 = { version = "0.5.1" } snap = { version = "1.0", optional = true } tokio-tungstenite = { version = "0.28", optional = true } zstd = { version = "0.13", optional = true } @@ -149,6 +149,7 @@ default_categorical_8 = ["minarrow/default_categorical_8"] extended_categorical = ["default_categorical_8", "minarrow/extended_categorical"] large_string = ["minarrow/large_string"] datetime = ["minarrow/datetime"] +decimal = ["minarrow/decimal"] extended_numeric_types = ["minarrow/extended_numeric_types"] lbuffer = ["minarrow/lbuffer"] # Schema-level Table metadata as key-value pairs diff --git a/rust/src/models/readers/ipc/window.rs b/rust/src/models/readers/ipc/window.rs index 4a09213..2dac296 100644 --- a/rust/src/models/readers/ipc/window.rs +++ b/rust/src/models/readers/ipc/window.rs @@ -91,6 +91,33 @@ fn window_array(array: &Array, offset: usize, len: usize) -> io::Result { NumericArray::Int16(arr) => win_int!(Int16, arr), #[cfg(feature = "extended_numeric_types")] NumericArray::UInt16(arr) => win_int!(UInt16, arr), + #[cfg(feature = "decimal")] + NumericArray::Decimal32(arr) => { + NumericArray::Decimal32(Arc::new(minarrow::DecimalArray { + data: window_buffer(&arr.data, offset, len), + null_mask: window_mask(arr.null_mask.as_ref(), offset, len), + precision: arr.precision, + scale: arr.scale, + })) + } + #[cfg(feature = "decimal")] + NumericArray::Decimal64(arr) => { + NumericArray::Decimal64(Arc::new(minarrow::DecimalArray { + data: window_buffer(&arr.data, offset, len), + null_mask: window_mask(arr.null_mask.as_ref(), offset, len), + precision: arr.precision, + scale: arr.scale, + })) + } + #[cfg(feature = "decimal")] + NumericArray::Decimal128(arr) => { + NumericArray::Decimal128(Arc::new(minarrow::DecimalArray { + data: window_buffer(&arr.data, offset, len), + null_mask: window_mask(arr.null_mask.as_ref(), offset, len), + precision: arr.precision, + scale: arr.scale, + })) + } NumericArray::Null => NumericArray::Null, }) } diff --git a/rust/src/models/types/parquet.rs b/rust/src/models/types/parquet.rs index 89f6a0d..b95782a 100644 --- a/rust/src/models/types/parquet.rs +++ b/rust/src/models/types/parquet.rs @@ -362,6 +362,12 @@ pub(crate) fn arrow_type_to_parquet( ArrowType::Duration64(_) => panic!("Duration does not map to a parquet type."), #[cfg(feature = "datetime")] ArrowType::Interval(_) => panic!("Interval does not map to a parquet type."), + #[cfg(feature = "decimal")] + ArrowType::Decimal32(_, _) + | ArrowType::Decimal64(_, _) + | ArrowType::Decimal128(_, _) => { + Err(IoError::UnsupportedType(format!("{ty:?}"))) + } #[cfg(all(feature = "extended_categorical", feature = "extended_numeric_types"))] &minarrow::ArrowType::Dictionary( minarrow::ffi::arrow_dtype::CategoricalIndexType::UInt16, From 05a9cdce50a73dc1c5934b9bdd730f5b29d6d979 Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Thu, 10 Sep 2026 04:29:42 +0100 Subject: [PATCH 2/7] Decimal type support for Arrow IPC and Parquet (TSK469) Arrow IPC: The FlatBuffers schema, record batch encoder, and decoder now handle Decimal32, Decimal64, and Decimal128 columns. The Decimal table struct carries precision, scale, and bitWidth. Data buffers are read and written as raw fixed-width integers matching the underlying storage type. Parquet: The type mapping, writer, and reader support decimal columns. Decimal32 and Decimal64 use INT32 and INT64 physical types with DECIMAL converted type. Decimal128 uses FIXED_LEN_BYTE_ARRAY(16) with big-endian two's complement encoding per the Parquet specification. All code is gated behind cfg(feature = "decimal"). --- rust/src/arrow/file.rs | 160 +++- rust/src/arrow/message.rs | 160 +++- rust/src/arrow/schema.rs | 160 +++- rust/src/models/decoders/ipc/parser.rs | 126 +++ rust/src/models/encoders/ipc/record_batch.rs | 12 + rust/src/models/encoders/ipc/schema.rs | 72 ++ rust/src/models/readers/parquet.rs | 835 ++++++++++++------- rust/src/models/types/parquet.rs | 47 +- rust/src/models/writers/parquet.rs | 123 ++- 9 files changed, 1374 insertions(+), 321 deletions(-) diff --git a/rust/src/arrow/file.rs b/rust/src/arrow/file.rs index a3e8713..77700b0 100644 --- a/rust/src/arrow/file.rs +++ b/rust/src/arrow/file.rs @@ -632,13 +632,14 @@ pub const ENUM_MIN_TYPE: u8 = 0; pub const ENUM_MAX_TYPE: u8 = 23; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_TYPE: [Type; 13] = [ +pub const ENUM_VALUES_TYPE: [Type; 14] = [ Type::NONE, Type::Null, Type::Int, Type::FloatingPoint, Type::Utf8, Type::Bool, + Type::Decimal, Type::Date, Type::Time, Type::Timestamp, @@ -666,6 +667,7 @@ impl Type { pub const FloatingPoint: Self = Self(3); pub const Utf8: Self = Self(5); pub const Bool: Self = Self(6); + pub const Decimal: Self = Self(7); pub const Date: Self = Self(8); pub const Time: Self = Self(9); pub const Timestamp: Self = Self(10); @@ -683,6 +685,7 @@ impl Type { Self::FloatingPoint, Self::Utf8, Self::Bool, + Self::Decimal, Self::Date, Self::Time, Self::Timestamp, @@ -700,6 +703,7 @@ impl Type { Self::FloatingPoint => Some("FloatingPoint"), Self::Utf8 => Some("Utf8"), Self::Bool => Some("Bool"), + Self::Decimal => Some("Decimal"), Self::Date => Some("Date"), Self::Time => Some("Time"), Self::Timestamp => Some("Timestamp"), @@ -1424,6 +1428,137 @@ impl core::fmt::Debug for Int<'_> { ds.finish() } } +pub enum DecimalOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct Decimal<'a> { + pub _tab: flatbuffers::Table<'a>, +} + +impl<'a> flatbuffers::Follow<'a> for Decimal<'a> { + type Inner = Decimal<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: flatbuffers::Table::new(buf, loc) } + } +} + +impl<'a> Decimal<'a> { + pub const VT_PRECISION: flatbuffers::VOffsetT = 4; + pub const VT_SCALE: flatbuffers::VOffsetT = 6; + pub const VT_BITWIDTH: flatbuffers::VOffsetT = 8; + + #[inline] + pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self { + Decimal { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>( + _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>, + args: &'args DecimalArgs + ) -> flatbuffers::WIPOffset> { + let mut builder = DecimalBuilder::new(_fbb); + builder.add_bitWidth(args.bitWidth); + builder.add_scale(args.scale); + builder.add_precision(args.precision); + builder.finish() + } + + + #[inline] + pub fn precision(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_PRECISION, Some(0)).unwrap()} + } + #[inline] + pub fn scale(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_SCALE, Some(0)).unwrap()} + } + #[inline] + pub fn bitWidth(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_BITWIDTH, Some(128)).unwrap()} + } +} + +impl flatbuffers::Verifiable for Decimal<'_> { + #[inline] + fn run_verifier( + v: &mut flatbuffers::Verifier, pos: usize + ) -> Result<(), flatbuffers::InvalidFlatbuffer> { + use self::flatbuffers::Verifiable; + v.visit_table(pos)? + .visit_field::("precision", Self::VT_PRECISION, false)? + .visit_field::("scale", Self::VT_SCALE, false)? + .visit_field::("bitWidth", Self::VT_BITWIDTH, false)? + .finish(); + Ok(()) + } +} +pub struct DecimalArgs { + pub precision: i32, + pub scale: i32, + pub bitWidth: i32, +} +impl<'a> Default for DecimalArgs { + #[inline] + fn default() -> Self { + DecimalArgs { + precision: 0, + scale: 0, + bitWidth: 128, + } + } +} + +pub struct DecimalBuilder<'a: 'b, 'b> { + fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>, + start_: flatbuffers::WIPOffset, +} +impl<'a: 'b, 'b> DecimalBuilder<'a, 'b> { + #[inline] + pub fn add_precision(&mut self, precision: i32) { + self.fbb_.push_slot::(Decimal::VT_PRECISION, precision, 0); + } + #[inline] + pub fn add_scale(&mut self, scale: i32) { + self.fbb_.push_slot::(Decimal::VT_SCALE, scale, 0); + } + #[inline] + pub fn add_bitWidth(&mut self, bitWidth: i32) { + self.fbb_.push_slot::(Decimal::VT_BITWIDTH, bitWidth, 128); + } + #[inline] + pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> DecimalBuilder<'a, 'b> { + let start = _fbb.start_table(); + DecimalBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + flatbuffers::WIPOffset::new(o.value()) + } +} + +impl core::fmt::Debug for Decimal<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut ds = f.debug_struct("Decimal"); + ds.field("precision", &self.precision()); + ds.field("scale", &self.scale()); + ds.field("bitWidth", &self.bitWidth()); + ds.finish() + } +} pub enum FloatingPointOffset {} #[derive(Copy, Clone, PartialEq)] @@ -2957,6 +3092,21 @@ impl<'a> Field<'a> { } } + #[inline] + #[allow(non_snake_case)] + pub fn type__as_decimal(&self) -> Option> { + if self.type_type() == Type::Decimal { + self.type_().map(|t| { + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + unsafe { Decimal::init_from_table(t) } + }) + } else { + None + } + } + #[inline] #[allow(non_snake_case)] pub fn type__as_date(&self) -> Option> { @@ -3080,6 +3230,7 @@ impl flatbuffers::Verifiable for Field<'_> { Type::FloatingPoint => v.verify_union_variant::>("Type::FloatingPoint", pos), Type::Utf8 => v.verify_union_variant::>("Type::Utf8", pos), Type::Bool => v.verify_union_variant::>("Type::Bool", pos), + Type::Decimal => v.verify_union_variant::>("Type::Decimal", pos), Type::Date => v.verify_union_variant::>("Type::Date", pos), Type::Time => v.verify_union_variant::>("Type::Time", pos), Type::Timestamp => v.verify_union_variant::>("Type::Timestamp", pos), @@ -3211,6 +3362,13 @@ impl core::fmt::Debug for Field<'_> { ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") } }, + Type::Decimal => { + if let Some(x) = self.type__as_decimal() { + ds.field("type_", &x) + } else { + ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") + } + }, Type::Date => { if let Some(x) = self.type__as_date() { ds.field("type_", &x) diff --git a/rust/src/arrow/message.rs b/rust/src/arrow/message.rs index 337dd30..c97cc3d 100644 --- a/rust/src/arrow/message.rs +++ b/rust/src/arrow/message.rs @@ -632,13 +632,14 @@ pub const ENUM_MIN_TYPE: u8 = 0; pub const ENUM_MAX_TYPE: u8 = 23; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_TYPE: [Type; 13] = [ +pub const ENUM_VALUES_TYPE: [Type; 14] = [ Type::NONE, Type::Null, Type::Int, Type::FloatingPoint, Type::Utf8, Type::Bool, + Type::Decimal, Type::Date, Type::Time, Type::Timestamp, @@ -666,6 +667,7 @@ impl Type { pub const FloatingPoint: Self = Self(3); pub const Utf8: Self = Self(5); pub const Bool: Self = Self(6); + pub const Decimal: Self = Self(7); pub const Date: Self = Self(8); pub const Time: Self = Self(9); pub const Timestamp: Self = Self(10); @@ -683,6 +685,7 @@ impl Type { Self::FloatingPoint, Self::Utf8, Self::Bool, + Self::Decimal, Self::Date, Self::Time, Self::Timestamp, @@ -700,6 +703,7 @@ impl Type { Self::FloatingPoint => Some("FloatingPoint"), Self::Utf8 => Some("Utf8"), Self::Bool => Some("Bool"), + Self::Decimal => Some("Decimal"), Self::Date => Some("Date"), Self::Time => Some("Time"), Self::Timestamp => Some("Timestamp"), @@ -1684,6 +1688,137 @@ impl core::fmt::Debug for Int<'_> { ds.finish() } } +pub enum DecimalOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct Decimal<'a> { + pub _tab: flatbuffers::Table<'a>, +} + +impl<'a> flatbuffers::Follow<'a> for Decimal<'a> { + type Inner = Decimal<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: flatbuffers::Table::new(buf, loc) } + } +} + +impl<'a> Decimal<'a> { + pub const VT_PRECISION: flatbuffers::VOffsetT = 4; + pub const VT_SCALE: flatbuffers::VOffsetT = 6; + pub const VT_BITWIDTH: flatbuffers::VOffsetT = 8; + + #[inline] + pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self { + Decimal { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>( + _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>, + args: &'args DecimalArgs + ) -> flatbuffers::WIPOffset> { + let mut builder = DecimalBuilder::new(_fbb); + builder.add_bitWidth(args.bitWidth); + builder.add_scale(args.scale); + builder.add_precision(args.precision); + builder.finish() + } + + + #[inline] + pub fn precision(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_PRECISION, Some(0)).unwrap()} + } + #[inline] + pub fn scale(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_SCALE, Some(0)).unwrap()} + } + #[inline] + pub fn bitWidth(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_BITWIDTH, Some(128)).unwrap()} + } +} + +impl flatbuffers::Verifiable for Decimal<'_> { + #[inline] + fn run_verifier( + v: &mut flatbuffers::Verifier, pos: usize + ) -> Result<(), flatbuffers::InvalidFlatbuffer> { + use self::flatbuffers::Verifiable; + v.visit_table(pos)? + .visit_field::("precision", Self::VT_PRECISION, false)? + .visit_field::("scale", Self::VT_SCALE, false)? + .visit_field::("bitWidth", Self::VT_BITWIDTH, false)? + .finish(); + Ok(()) + } +} +pub struct DecimalArgs { + pub precision: i32, + pub scale: i32, + pub bitWidth: i32, +} +impl<'a> Default for DecimalArgs { + #[inline] + fn default() -> Self { + DecimalArgs { + precision: 0, + scale: 0, + bitWidth: 128, + } + } +} + +pub struct DecimalBuilder<'a: 'b, 'b> { + fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>, + start_: flatbuffers::WIPOffset, +} +impl<'a: 'b, 'b> DecimalBuilder<'a, 'b> { + #[inline] + pub fn add_precision(&mut self, precision: i32) { + self.fbb_.push_slot::(Decimal::VT_PRECISION, precision, 0); + } + #[inline] + pub fn add_scale(&mut self, scale: i32) { + self.fbb_.push_slot::(Decimal::VT_SCALE, scale, 0); + } + #[inline] + pub fn add_bitWidth(&mut self, bitWidth: i32) { + self.fbb_.push_slot::(Decimal::VT_BITWIDTH, bitWidth, 128); + } + #[inline] + pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> DecimalBuilder<'a, 'b> { + let start = _fbb.start_table(); + DecimalBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + flatbuffers::WIPOffset::new(o.value()) + } +} + +impl core::fmt::Debug for Decimal<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut ds = f.debug_struct("Decimal"); + ds.field("precision", &self.precision()); + ds.field("scale", &self.scale()); + ds.field("bitWidth", &self.bitWidth()); + ds.finish() + } +} pub enum FloatingPointOffset {} #[derive(Copy, Clone, PartialEq)] @@ -3217,6 +3352,21 @@ impl<'a> Field<'a> { } } + #[inline] + #[allow(non_snake_case)] + pub fn type__as_decimal(&self) -> Option> { + if self.type_type() == Type::Decimal { + self.type_().map(|t| { + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + unsafe { Decimal::init_from_table(t) } + }) + } else { + None + } + } + #[inline] #[allow(non_snake_case)] pub fn type__as_date(&self) -> Option> { @@ -3340,6 +3490,7 @@ impl flatbuffers::Verifiable for Field<'_> { Type::FloatingPoint => v.verify_union_variant::>("Type::FloatingPoint", pos), Type::Utf8 => v.verify_union_variant::>("Type::Utf8", pos), Type::Bool => v.verify_union_variant::>("Type::Bool", pos), + Type::Decimal => v.verify_union_variant::>("Type::Decimal", pos), Type::Date => v.verify_union_variant::>("Type::Date", pos), Type::Time => v.verify_union_variant::>("Type::Time", pos), Type::Timestamp => v.verify_union_variant::>("Type::Timestamp", pos), @@ -3471,6 +3622,13 @@ impl core::fmt::Debug for Field<'_> { ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") } }, + Type::Decimal => { + if let Some(x) = self.type__as_decimal() { + ds.field("type_", &x) + } else { + ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") + } + }, Type::Date => { if let Some(x) = self.type__as_date() { ds.field("type_", &x) diff --git a/rust/src/arrow/schema.rs b/rust/src/arrow/schema.rs index a68a118..58e34d9 100644 --- a/rust/src/arrow/schema.rs +++ b/rust/src/arrow/schema.rs @@ -632,13 +632,14 @@ pub const ENUM_MIN_TYPE: u8 = 0; pub const ENUM_MAX_TYPE: u8 = 23; #[deprecated(since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021.")] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_TYPE: [Type; 13] = [ +pub const ENUM_VALUES_TYPE: [Type; 14] = [ Type::NONE, Type::Null, Type::Int, Type::FloatingPoint, Type::Utf8, Type::Bool, + Type::Decimal, Type::Date, Type::Time, Type::Timestamp, @@ -666,6 +667,7 @@ impl Type { pub const FloatingPoint: Self = Self(3); pub const Utf8: Self = Self(5); pub const Bool: Self = Self(6); + pub const Decimal: Self = Self(7); pub const Date: Self = Self(8); pub const Time: Self = Self(9); pub const Timestamp: Self = Self(10); @@ -683,6 +685,7 @@ impl Type { Self::FloatingPoint, Self::Utf8, Self::Bool, + Self::Decimal, Self::Date, Self::Time, Self::Timestamp, @@ -700,6 +703,7 @@ impl Type { Self::FloatingPoint => Some("FloatingPoint"), Self::Utf8 => Some("Utf8"), Self::Bool => Some("Bool"), + Self::Decimal => Some("Decimal"), Self::Date => Some("Date"), Self::Time => Some("Time"), Self::Timestamp => Some("Timestamp"), @@ -1264,6 +1268,137 @@ impl core::fmt::Debug for Int<'_> { ds.finish() } } +pub enum DecimalOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct Decimal<'a> { + pub _tab: flatbuffers::Table<'a>, +} + +impl<'a> flatbuffers::Follow<'a> for Decimal<'a> { + type Inner = Decimal<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: flatbuffers::Table::new(buf, loc) } + } +} + +impl<'a> Decimal<'a> { + pub const VT_PRECISION: flatbuffers::VOffsetT = 4; + pub const VT_SCALE: flatbuffers::VOffsetT = 6; + pub const VT_BITWIDTH: flatbuffers::VOffsetT = 8; + + #[inline] + pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self { + Decimal { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr>( + _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr>, + args: &'args DecimalArgs + ) -> flatbuffers::WIPOffset> { + let mut builder = DecimalBuilder::new(_fbb); + builder.add_bitWidth(args.bitWidth); + builder.add_scale(args.scale); + builder.add_precision(args.precision); + builder.finish() + } + + + #[inline] + pub fn precision(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_PRECISION, Some(0)).unwrap()} + } + #[inline] + pub fn scale(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_SCALE, Some(0)).unwrap()} + } + #[inline] + pub fn bitWidth(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(Decimal::VT_BITWIDTH, Some(128)).unwrap()} + } +} + +impl flatbuffers::Verifiable for Decimal<'_> { + #[inline] + fn run_verifier( + v: &mut flatbuffers::Verifier, pos: usize + ) -> Result<(), flatbuffers::InvalidFlatbuffer> { + use self::flatbuffers::Verifiable; + v.visit_table(pos)? + .visit_field::("precision", Self::VT_PRECISION, false)? + .visit_field::("scale", Self::VT_SCALE, false)? + .visit_field::("bitWidth", Self::VT_BITWIDTH, false)? + .finish(); + Ok(()) + } +} +pub struct DecimalArgs { + pub precision: i32, + pub scale: i32, + pub bitWidth: i32, +} +impl<'a> Default for DecimalArgs { + #[inline] + fn default() -> Self { + DecimalArgs { + precision: 0, + scale: 0, + bitWidth: 128, + } + } +} + +pub struct DecimalBuilder<'a: 'b, 'b> { + fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a>, + start_: flatbuffers::WIPOffset, +} +impl<'a: 'b, 'b> DecimalBuilder<'a, 'b> { + #[inline] + pub fn add_precision(&mut self, precision: i32) { + self.fbb_.push_slot::(Decimal::VT_PRECISION, precision, 0); + } + #[inline] + pub fn add_scale(&mut self, scale: i32) { + self.fbb_.push_slot::(Decimal::VT_SCALE, scale, 0); + } + #[inline] + pub fn add_bitWidth(&mut self, bitWidth: i32) { + self.fbb_.push_slot::(Decimal::VT_BITWIDTH, bitWidth, 128); + } + #[inline] + pub fn new(_fbb: &'b mut flatbuffers::FlatBufferBuilder<'a>) -> DecimalBuilder<'a, 'b> { + let start = _fbb.start_table(); + DecimalBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + flatbuffers::WIPOffset::new(o.value()) + } +} + +impl core::fmt::Debug for Decimal<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut ds = f.debug_struct("Decimal"); + ds.field("precision", &self.precision()); + ds.field("scale", &self.scale()); + ds.field("bitWidth", &self.bitWidth()); + ds.finish() + } +} pub enum FloatingPointOffset {} #[derive(Copy, Clone, PartialEq)] @@ -2797,6 +2932,21 @@ impl<'a> Field<'a> { } } + #[inline] + #[allow(non_snake_case)] + pub fn type__as_decimal(&self) -> Option> { + if self.type_type() == Type::Decimal { + self.type_().map(|t| { + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + unsafe { Decimal::init_from_table(t) } + }) + } else { + None + } + } + #[inline] #[allow(non_snake_case)] pub fn type__as_date(&self) -> Option> { @@ -2920,6 +3070,7 @@ impl flatbuffers::Verifiable for Field<'_> { Type::FloatingPoint => v.verify_union_variant::>("Type::FloatingPoint", pos), Type::Utf8 => v.verify_union_variant::>("Type::Utf8", pos), Type::Bool => v.verify_union_variant::>("Type::Bool", pos), + Type::Decimal => v.verify_union_variant::>("Type::Decimal", pos), Type::Date => v.verify_union_variant::>("Type::Date", pos), Type::Time => v.verify_union_variant::>("Type::Time", pos), Type::Timestamp => v.verify_union_variant::>("Type::Timestamp", pos), @@ -3051,6 +3202,13 @@ impl core::fmt::Debug for Field<'_> { ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") } }, + Type::Decimal => { + if let Some(x) = self.type__as_decimal() { + ds.field("type_", &x) + } else { + ds.field("type_", &"InvalidFlatbuffer: Union discriminant does not match value.") + } + }, Type::Date => { if let Some(x) = self.type__as_date() { ds.field("type_", &x) diff --git a/rust/src/models/decoders/ipc/parser.rs b/rust/src/models/decoders/ipc/parser.rs index 8868304..fdc0f78 100644 --- a/rust/src/models/decoders/ipc/parser.rs +++ b/rust/src/models/decoders/ipc/parser.rs @@ -38,6 +38,8 @@ use log::warn; use tracing::debug; use flatbuffers::Vector; +#[cfg(feature = "decimal")] +use minarrow::DecimalArray; #[cfg(feature = "datetime")] use minarrow::enums::time_units::TimeUnit as MnTimeUnit; use minarrow::ffi::arrow_dtype::{ArrowType, CategoricalIndexType}; @@ -380,6 +382,47 @@ impl RecordBatchParser { data, null_mask, )))) } + // decimal + #[cfg(feature = "decimal")] + ArrowType::Decimal32(precision, scale) => { + let (slice, _) = Self::extract_buffer_slice( + &fbuf_meta, + &mut buffer_idx, + arrow_buf, + &field.name, + corrections, + )?; + let data = + unsafe { Self::buffer_from_slice::(slice, field_len, &arc_opt) }; + Array::from_decimal32(DecimalArray::new(data, null_mask, *precision, *scale)) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal64(precision, scale) => { + let (slice, _) = Self::extract_buffer_slice( + &fbuf_meta, + &mut buffer_idx, + arrow_buf, + &field.name, + corrections, + )?; + let data = + unsafe { Self::buffer_from_slice::(slice, field_len, &arc_opt) }; + Array::from_decimal64(DecimalArray::new(data, null_mask, *precision, *scale)) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal128(precision, scale) => { + let (slice, _) = Self::extract_buffer_slice( + &fbuf_meta, + &mut buffer_idx, + arrow_buf, + &field.name, + corrections, + )?; + let data = + unsafe { Self::buffer_from_slice::(slice, field_len, &arc_opt) }; + Array::from_decimal128(DecimalArray::new(data, null_mask, *precision, *scale)) + } + // dictionary ArrowType::Dictionary(idx_ty) => { // indices @@ -1108,6 +1151,61 @@ pub fn decode_record_batch( make_numeric_array(&field.dtype, data, null_mask)? } + #[cfg(feature = "decimal")] + ArrowType::Decimal32(precision, scale) => { + let (off, len) = consume_buffer( + &buffers, + &mut buffer_idx, + body_start, + body_len, + &field.name, + corrections, + )?; + let data = shared.slice(off..off + len); + Array::from_decimal32(DecimalArray::new( + minarrow::Buffer::from_shared(data), + null_mask, + *precision, + *scale, + )) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal64(precision, scale) => { + let (off, len) = consume_buffer( + &buffers, + &mut buffer_idx, + body_start, + body_len, + &field.name, + corrections, + )?; + let data = shared.slice(off..off + len); + Array::from_decimal64(DecimalArray::new( + minarrow::Buffer::from_shared(data), + null_mask, + *precision, + *scale, + )) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal128(precision, scale) => { + let (off, len) = consume_buffer( + &buffers, + &mut buffer_idx, + body_start, + body_len, + &field.name, + corrections, + )?; + let data = shared.slice(off..off + len); + Array::from_decimal128(DecimalArray::new( + minarrow::Buffer::from_shared(data), + null_mask, + *precision, + *scale, + )) + } + ArrowType::Dictionary(_idx_ty) => { let dict_key = col_idx as i64; let dict_values = dicts.get(&dict_key).ok_or_else(|| { @@ -1167,6 +1265,8 @@ fn data_buffer_count(dtype: &ArrowType) -> usize { ArrowType::String => 2, #[cfg(feature = "large_string")] ArrowType::LargeString => 2, + #[cfg(feature = "decimal")] + ArrowType::Decimal32(_, _) | ArrowType::Decimal64(_, _) | ArrowType::Decimal128(_, _) => 1, _ => 1, } } @@ -1761,6 +1861,19 @@ fn extract_base_type(fb_field: &fb::Field) -> io::Result { })?; Ok(ArrowType::Duration64(convert_time_unit_fb(d.unit())?)) } + #[cfg(feature = "decimal")] + fb::Type::Decimal => { + let d = fb_field + .type__as_decimal() + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing Decimal type"))?; + let precision = d.precision() as u8; + let scale = d.scale() as i8; + match d.bitWidth() { + 32 => Ok(ArrowType::Decimal32(precision, scale)), + 64 => Ok(ArrowType::Decimal64(precision, scale)), + _ => Ok(ArrowType::Decimal128(precision, scale)), + } + } fb::Type::Bool => Ok(ArrowType::Boolean), other => { if let Some(dict) = fb_field.dictionary() { @@ -1963,6 +2076,19 @@ pub fn convert_fb_field_to_arrow( })?; ArrowType::Duration64(convert_time_unit_fbf(d.unit())?) } + #[cfg(feature = "decimal")] + fbf::Type::Decimal => { + let d = fbf_field.type__as_decimal().ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "missing Decimal type") + })?; + let precision = d.precision() as u8; + let scale = d.scale() as i8; + match d.bitWidth() { + 32 => ArrowType::Decimal32(precision, scale), + 64 => ArrowType::Decimal64(precision, scale), + _ => ArrowType::Decimal128(precision, scale), + } + } other => { return Err(io::Error::new( io::ErrorKind::InvalidData, diff --git a/rust/src/models/encoders/ipc/record_batch.rs b/rust/src/models/encoders/ipc/record_batch.rs index d2a262e..7d46f8a 100644 --- a/rust/src/models/encoders/ipc/record_batch.rs +++ b/rust/src/models/encoders/ipc/record_batch.rs @@ -224,6 +224,18 @@ pub(crate) fn compute_body_layout<'a, B: StreamBuffer>( NumericArray::UInt16(arr) => { (as_bytes(&arr.data.as_slice()[offset..offset + len]), arr.null_mask.as_ref()) } + #[cfg(feature = "decimal")] + NumericArray::Decimal32(arr) => { + (as_bytes(&arr.data.as_slice()[offset..offset + len]), arr.null_mask.as_ref()) + } + #[cfg(feature = "decimal")] + NumericArray::Decimal64(arr) => { + (as_bytes(&arr.data.as_slice()[offset..offset + len]), arr.null_mask.as_ref()) + } + #[cfg(feature = "decimal")] + NumericArray::Decimal128(arr) => { + (as_bytes(&arr.data.as_slice()[offset..offset + len]), arr.null_mask.as_ref()) + } _ => { return Err(io::Error::new( io::ErrorKind::InvalidInput, diff --git a/rust/src/models/encoders/ipc/schema.rs b/rust/src/models/encoders/ipc/schema.rs index 84bfe3d..6d541cc 100644 --- a/rust/src/models/encoders/ipc/schema.rs +++ b/rust/src/models/encoders/ipc/schema.rs @@ -317,6 +317,42 @@ fn build_flatbuf_field<'fbb>( let duration = fbm::Duration::create(fbb, &fbm::DurationArgs { unit }); (fbm::Type::Duration, Some(duration.as_union_value()), None) } + #[cfg(feature = "decimal")] + ArrowType::Decimal32(precision, scale) => { + let decimal = fbm::Decimal::create( + fbb, + &fbm::DecimalArgs { + precision: *precision as i32, + scale: *scale as i32, + bitWidth: 32, + }, + ); + (fbm::Type::Decimal, Some(decimal.as_union_value()), None) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal64(precision, scale) => { + let decimal = fbm::Decimal::create( + fbb, + &fbm::DecimalArgs { + precision: *precision as i32, + scale: *scale as i32, + bitWidth: 64, + }, + ); + (fbm::Type::Decimal, Some(decimal.as_union_value()), None) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal128(precision, scale) => { + let decimal = fbm::Decimal::create( + fbb, + &fbm::DecimalArgs { + precision: *precision as i32, + scale: *scale as i32, + bitWidth: 128, + }, + ); + (fbm::Type::Decimal, Some(decimal.as_union_value()), None) + } ArrowType::Dictionary(idx_ty) => { // Build index type for dictionary let idx_width = match idx_ty { @@ -810,6 +846,42 @@ fn build_flatbuf_field_file<'fbb>( let duration = fbf::Duration::create(fbb, &fbf::DurationArgs { unit }); (fbf::Type::Duration, Some(duration.as_union_value()), None) } + #[cfg(feature = "decimal")] + ArrowType::Decimal32(precision, scale) => { + let decimal = fbf::Decimal::create( + fbb, + &fbf::DecimalArgs { + precision: *precision as i32, + scale: *scale as i32, + bitWidth: 32, + }, + ); + (fbf::Type::Decimal, Some(decimal.as_union_value()), None) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal64(precision, scale) => { + let decimal = fbf::Decimal::create( + fbb, + &fbf::DecimalArgs { + precision: *precision as i32, + scale: *scale as i32, + bitWidth: 64, + }, + ); + (fbf::Type::Decimal, Some(decimal.as_union_value()), None) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal128(precision, scale) => { + let decimal = fbf::Decimal::create( + fbb, + &fbf::DecimalArgs { + precision: *precision as i32, + scale: *scale as i32, + bitWidth: 128, + }, + ); + (fbf::Type::Decimal, Some(decimal.as_union_value()), None) + } ArrowType::Dictionary(idx_ty) => { let idx_width = match idx_ty { #[cfg(any( diff --git a/rust/src/models/readers/parquet.rs b/rust/src/models/readers/parquet.rs index 4e8bc42..558639f 100644 --- a/rust/src/models/readers/parquet.rs +++ b/rust/src/models/readers/parquet.rs @@ -6,10 +6,15 @@ //! # Parquet table reader - *reads into `minarrow::Table`* //! +//! Reads flat-schema Parquet files written by lightstream or by other +//! writers such as parquet-cpp (pyarrow) into a fully materialised Table. +//! //! ## Features -//! - Supports DataPageV2 plus the legacy V1 Parquet layout -//! - Decodes hybrid RLE/bit-packed definition levels and dictionary indices -//! - Handles `PLAIN` and `RLE_DICTIONARY` value encodings +//! - DataPageV1 and DataPageV2 page layouts, across any number of row groups +//! - Hybrid RLE/bit-packed definition levels and dictionary indices +//! - `PLAIN`, `RLE_DICTIONARY` and `PLAIN_DICTIONARY` value encodings, plus +//! `RLE` booleans. Dictionary-encoded columns of any physical type are +//! expanded on read. //! - Optional feature-gated Snappy / Zstd compression //! - Type maps to Arrow/Minarrow - {i32, i64, u32, u64, f32, f64, bool, utf8 //! dictionary Option { + #[cfg(feature = "decimal")] + if se.converted_type == Some(5) { + let precision = se.precision.unwrap_or(0) as u8; + let scale = se.scale.unwrap_or(0) as i8; + return Some(ParquetLogicalType::DecimalType { precision, scale }); + } + ParquetLogicalType::from_converted_type(se.converted_type) +} + +/// Decode a FIXED_LEN_BYTE_ARRAY Decimal128 buffer (16 big-endian bytes per +/// value) into a `Vec64`. +#[cfg(feature = "decimal")] +fn decode_decimal128_plain(buf: &[u8]) -> Result, IoError> { + if buf.len() % 16 != 0 { + return Err(IoError::Format( + "decode_decimal128_plain: buffer len % 16 != 0".into(), + )); + } + Ok(buf + .chunks_exact(16) + .map(|c| i128::from_be_bytes(c.try_into().unwrap())) + .collect()) +} + /// Read an entire in-memory Table from a Parquet v2 file. /// // TODO: Serial read throughput on the chunked bench is ~390 MiB/s @@ -96,17 +139,36 @@ pub fn load_parquet_table_cols( r: R, columns: &[&str], ) -> Result { - let projection: std::collections::HashSet = + let projection: HashSet = columns.iter().map(|s| s.to_string()).collect(); read_parquet_impl(r, Some(projection)) } /// Shared implementation for full and projected Parquet reads. -/// When `projection` is `Some`, only columns whose names appear in the set -/// are read. Skipped columns never hit disk. +/// +/// Every row group contributes its column chunks to the same output +/// columns, so a file with several row groups reads as one Table. When +/// `projection` is `Some`, only columns whose names appear in the set are +/// read. Skipped columns never hit disk. +/// +/// ## Behaviour +/// - Values arrive in the Parquet layout, where a page's value section +/// holds non-null entries only. Each page is decoded to PLAIN entries and +/// scattered against its definition levels, so [`decode_column`] always +/// receives one entry per row. +/// - Dictionary-encoded pages of any physical type are expanded through +/// the row group's dictionary, so dictionary encoding stays a storage +/// detail of the writer. +/// - A BYTE_ARRAY UTF8 leaf with a dictionary page in a file carrying +/// [`PARQUET_CREATED_BY`](crate::constants::PARQUET_CREATED_BY) is the +/// lightstream categorical convention and reads back as +/// `ArrowType::Dictionary`. Dictionaries from several row groups merge +/// into one set of unique values. +/// - Repeated (nested) columns and unknown compression codecs are +/// reported as errors rather than decoded. fn read_parquet_impl( mut r: R, - projection: Option>, + projection: Option>, ) -> Result { // read the 8-byte footer r.seek(SeekFrom::End(-8))?; @@ -121,13 +183,22 @@ fn read_parquet_impl( r.seek(SeekFrom::End(-8 - footer_len as i64))?; let mut footer = vec![0u8; footer_len as usize]; r.read_exact(&mut footer)?; - let mut cur = std::io::Cursor::new(&footer); + let mut cur = Cursor::new(&footer); let meta = parse_file_metadata(&mut cur)?; + // Leaf schema elements carry a physical type and the root group element + // does not. Leaves sit in schema order, which is also the order of the + // column chunks inside every row group. + let leaves: Vec<(&SchemaElement, ParquetPhysicalType)> = meta + .schema + .iter() + .filter_map(|se| se.type_.map(|ty| (se, ty))) + .collect(); + // Validate projection names against schema before reading any data if let Some(ref proj) = projection { for name in proj { - if !meta.schema.iter().any(|se| se.name == *name) { + if !leaves.iter().any(|(se, _)| se.name == *name) { return Err(IoError::Format(format!( "column '{}' not found in schema", name @@ -136,117 +207,226 @@ fn read_parquet_impl( } } - // map Parquet schema -> Arrow types. The schema list opens with a root - // group element carrying no physical type; leaf column elements follow. - // Map only the leaves, in schema order, so indices line up with the row - // group's column chunks. - let arrow_types: Vec<_> = meta - .schema + if meta.row_groups.is_empty() { + return Err(IoError::Format("file has no row groups".into())); + } + if let Some(rg) = meta + .row_groups .iter() - .filter_map(|se| se.type_.map(|ty| (ty, se.converted_type))) - .map(|(ty, converted)| { - parquet_to_arrow_type(ty, ParquetLogicalType::from_converted_type(converted)) - }) - .collect::>()?; + .find(|rg| rg.columns.len() != leaves.len()) + { + return Err(IoError::Format(format!( + "row group has {} column chunks for {} schema leaves", + rg.columns.len(), + leaves.len() + ))); + } - // single row-group, flat schema only - let rg = &meta.row_groups[0]; - let mut columns = Vec::with_capacity(rg.columns.len()); + // The categorical convention only applies to lightstream's own files. + // Other writers dictionary-encode any column type as a storage detail. + let lightstream_written = meta.created_by.as_deref() == Some(PARQUET_CREATED_BY); - for (col_idx, chunk) in rg.columns.iter().enumerate() { - let cmeta = &chunk.meta_data; - let col_name = &cmeta.path_in_schema[0]; + let mut columns = Vec::with_capacity(leaves.len()); + for (col_idx, &(leaf, physical)) in leaves.iter().enumerate() { // Skip columns not in the projection if let Some(ref proj) = projection - && !proj.contains(col_name) + && !proj.contains(&leaf.name) { continue; } - // The Parquet schema has no slot to record that a column was - // originally a Dictionary - the writer maps Dictionary(_) to a - // physical Int32 with no logical type. Recover the Arrow shape - // from the column metadata: a column carrying a dictionary page - // is decoded as Dictionary using the build's default index - // width (UInt8 when `default_categorical_8` is on, UInt32 - // otherwise). Round-trip across feature flags isn't supported. - let schema_ty = &arrow_types[col_idx]; - let dictionary_ty = if cmeta.dictionary_page_offset.is_some() { - Some(ArrowType::Dictionary(default_categorical_index_type())) + // OPTIONAL leaves carry one definition level per row. REQUIRED + // leaves carry none. REPEATED leaves describe nested data. + let max_def_level: u8 = match leaf.repetition_type { + 0 => 0, + 1 => 1, + other => { + return Err(IoError::UnsupportedType(format!( + "repeated column '{}' (repetition_type {other})", + leaf.name + ))); + } + }; + + let logical = logical_type_from_schema(leaf); + let schema_ty = parquet_to_arrow_type(physical, logical)?; + let has_dictionary = meta + .row_groups + .iter() + .any(|rg| rg.columns[col_idx].meta_data.dictionary_page_offset.is_some()); + let categorical = + lightstream_written && physical == ParquetPhysicalType::ByteArray && has_dictionary; + let ty = if categorical { + ArrowType::Dictionary(default_categorical_index_type()) } else { - None + schema_ty }; - let ty: &ArrowType = dictionary_ty.as_ref().unwrap_or(schema_ty); - - // read the DICTIONARY_PAGE if present - let dict = if let Some(dict_off) = cmeta.dictionary_page_offset { - r.seek(SeekFrom::Start(dict_off as u64))?; - let ph = parse_page_header(&mut r)?; - if ph.type_ != PageType::DictionaryPage { - return Err(IoError::Format("expected DICTIONARY_PAGE".into())); - } - let mut compr = vec![0u8; ph.compressed_page_size as usize]; - r.read_exact(&mut compr)?; - match map_codec(cmeta.codec) { - Some(c) => parse_dictionary_values(&decompress(&compr, c)?)?, - None => parse_dictionary_values(&compr)?, - } + // Categorical columns accumulate their dictionary indices as PLAIN + // u32 entries. Every other column accumulates its physical values. + let layout = if categorical { + ValueLayout::Fixed(4) } else { - Vec::new() + ValueLayout::of(physical, leaf.type_length)? }; - // walk all the DATA_PAGE_V2s in a row - let total_vals = cmeta.num_values as usize; - let mut def_levels = Vec::with_capacity(total_vals); - let mut values_buf = Vec::new(); - let mut pages_read = 0; - let mut page_encoding = ParquetEncoding::Plain; - - // seek once to the first data page - r.seek(SeekFrom::Start(cmeta.data_page_offset as u64))?; - - while pages_read < total_vals { - // parse the next page header - let ph = parse_page_header(&mut r)?; - let (page_defs, enc, page_vals) = match ph.type_ { - PageType::DataPageV2 => read_data_page_v2(&mut r, &ph, cmeta)?, - PageType::DataPage => read_data_page_v1(&mut r, &ph, cmeta)?, - t => return Err(IoError::Format(format!("unsupported page type {:?}", t))), - }; - if pages_read == 0 { - page_encoding = enc; + let mut def_levels: Vec = Vec::new(); + let mut values: Vec = Vec::new(); + // Merged categorical dictionary across row groups, with the lookup + // used to remap each row group's local indices onto it. + let mut unique_values: Vec> = Vec::new(); + let mut unique_index: BTreeMap, u32> = BTreeMap::new(); + + for rg in &meta.row_groups { + let cmeta = &rg.columns[col_idx].meta_data; + let codec = map_codec(cmeta.codec)?; + + // The row group's dictionary entries without their PLAIN length + // prefix, plus the remap onto the merged categorical dictionary. + let mut dict_entries: Vec> = Vec::new(); + let mut dict_remap: Vec = Vec::new(); + if let Some(dict_off) = cmeta.dictionary_page_offset { + r.seek(SeekFrom::Start(dict_off as u64))?; + let ph = parse_page_header(&mut r)?; + if ph.type_ != PageType::DictionaryPage { + return Err(IoError::Format("expected DICTIONARY_PAGE".into())); + } + let mut body = vec![0u8; ph.compressed_page_size as usize]; + r.read_exact(&mut body)?; + let body = match codec { + Some(c) => decompress(&body, c)?, + None => body, + }; + if categorical { + for entry in parse_dictionary_values(&body)? { + let merged = match unique_index.get(&entry) { + Some(&idx) => idx, + None => { + let idx = unique_values.len() as u32; + unique_index.insert(entry.clone(), idx); + unique_values.push(entry); + idx + } + }; + dict_remap.push(merged); + } + } else { + dict_entries = match layout { + ValueLayout::LengthPrefixed => parse_dictionary_values(&body)?, + ValueLayout::Fixed(width) => { + if body.len() % width != 0 { + return Err(IoError::Format(format!( + "dictionary page of {} bytes is not a multiple of the {width}-byte value width", + body.len() + ))); + } + body.chunks_exact(width).map(<[u8]>::to_vec).collect() + } + }; + } } - // accumulate `page_defs.len()` logical rows - let this_count = page_defs.len().min(total_vals - pages_read); - def_levels.extend_from_slice(&page_defs[..this_count]); - values_buf.extend_from_slice(&page_vals); - pages_read += this_count; + // Walk the chunk's data pages. num_values counts every row of + // the chunk, nulls included, which is the definition level + // count. + r.seek(SeekFrom::Start(cmeta.data_page_offset as u64))?; + let chunk_levels = cmeta.num_values as usize; + let mut levels_read = 0usize; + + while levels_read < chunk_levels { + let ph = parse_page_header(&mut r)?; + let (page_defs, encoding, page_values) = match ph.type_ { + PageType::DataPage => read_data_page_v1(&mut r, &ph, codec, max_def_level)?, + PageType::DataPageV2 => { + read_data_page_v2(&mut r, &ph, codec, max_def_level)? + } + t => return Err(IoError::Format(format!("unsupported page type {:?}", t))), + }; + let n_valid = page_defs.iter().filter(|&&v| v).count(); + + // PLAIN entries for the page's non-null values. Boolean pages + // unpack to one byte per value here so every physical type + // scatters through the same fixed or length-prefixed layout. + let plain: Vec = match (encoding, physical) { + (ParquetEncoding::Plain, ParquetPhysicalType::Boolean) => { + if page_values.len() * 8 < n_valid { + return Err(IoError::Format("truncated boolean page".into())); + } + (0..n_valid) + .map(|i| (page_values[i / 8] >> (i % 8)) & 1) + .collect() + } + (ParquetEncoding::Rle, ParquetPhysicalType::Boolean) => { + // RLE booleans carry a 4-byte run length ahead of + // the hybrid run, like V1 levels. + let run = length_prefixed_run(&page_values)?; + decode_hybrid(run, 1, n_valid)? + .iter() + .map(|&v| v as u8) + .collect() + } + (ParquetEncoding::Plain, _) => page_values, + (ParquetEncoding::RleDictionary | ParquetEncoding::PlainDictionary, _) => { + let indices = if n_valid == 0 { + Vec64::new() + } else { + decode_dictionary_indices_rle(&page_values, n_valid)? + }; + let mut plain = Vec::new(); + for &idx in indices.iter() { + if categorical { + let merged = dict_remap.get(idx as usize).ok_or_else(|| { + IoError::Format(format!("dictionary index {idx} out of range")) + })?; + plain.extend_from_slice(&merged.to_le_bytes()); + continue; + } + let entry = dict_entries.get(idx as usize).ok_or_else(|| { + IoError::Format(format!("dictionary index {idx} out of range")) + })?; + if layout == ValueLayout::LengthPrefixed { + plain.extend_from_slice(&(entry.len() as u32).to_le_bytes()); + } + plain.extend_from_slice(entry); + } + plain + } + (enc, _) => { + return Err(IoError::UnsupportedEncoding(format!( + "{enc:?} for column '{}'", + leaf.name + ))); + } + }; - // cursor is at the start of the next page header + scatter_by_definition_levels(&mut values, &plain, &page_defs, layout)?; + levels_read += page_defs.len(); + def_levels.extend(page_defs); + } + } + + if def_levels.len() != meta.num_rows as usize { + return Err(IoError::Format(format!( + "column '{}' holds {} rows, file metadata says {}", + leaf.name, + def_levels.len(), + meta.num_rows + ))); } - // decode the column array - let array = decode_column( - ty, - page_encoding, - &dict, - &values_buf, - total_vals, - def_levels.clone(), - )?; + let null_count = def_levels.iter().filter(|&&b| !b).count(); + let array = decode_column(&ty, &unique_values, &values, def_levels.len(), def_levels)?; columns.push(FieldArray { field: Field { - name: col_name.clone(), - dtype: ty.clone(), - nullable: chunk.meta_data.definition_level >= 1 || def_levels.iter().any(|&b| !b), + name: leaf.name.clone(), + dtype: ty, + nullable: max_def_level > 0, metadata: Default::default(), } .into(), array, - null_count: def_levels.iter().filter(|&&b| !b).count(), + null_count, }); } @@ -258,24 +438,165 @@ fn read_parquet_impl( }) } -/// DataPageV2 reader: read exactly `compressed_page_size` bytes, split into -/// rep / def, decompress the remainder, decode def‐levels, return the raw -/// values and the page encoding. +/// Byte layout of one value inside a PLAIN value section. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum ValueLayout { + /// Every value occupies the given number of bytes. + Fixed(usize), + /// A 4-byte little-endian length precedes each value's bytes. + LengthPrefixed, +} + +impl ValueLayout { + /// Layout of the physical type once a page has been decoded to PLAIN + /// entries. Booleans count as one byte per value because the page + /// readers unpack them before scattering. `type_length` is the schema + /// element's width for FIXED_LEN_BYTE_ARRAY leaves. + fn of(physical: ParquetPhysicalType, type_length: Option) -> Result { + Ok(match physical { + ParquetPhysicalType::Boolean => ValueLayout::Fixed(1), + ParquetPhysicalType::Int32 | ParquetPhysicalType::Float => ValueLayout::Fixed(4), + ParquetPhysicalType::Int64 | ParquetPhysicalType::Double => ValueLayout::Fixed(8), + ParquetPhysicalType::ByteArray => ValueLayout::LengthPrefixed, + ParquetPhysicalType::FixedLenByteArray => match type_length { + Some(width) if width > 0 => ValueLayout::Fixed(width as usize), + _ => { + return Err(IoError::Format( + "FIXED_LEN_BYTE_ARRAY leaf without a type_length".into(), + )); + } + }, + }) + } +} + +/// Append a page's PLAIN entries to `out` with a zero entry at every null +/// row, so the column buffer holds one entry per row. +/// +/// Parquet value sections hold non-null values only, while the PLAIN +/// decoders behind [`decode_column`] expect one entry per row with the null +/// mask applied afterwards. A null row becomes zero bytes for a fixed-width +/// value or a zero length for a length-prefixed value. +fn scatter_by_definition_levels( + out: &mut Vec, + plain: &[u8], + def_levels: &[bool], + layout: ValueLayout, +) -> Result<(), IoError> { + if def_levels.iter().all(|&v| v) { + out.extend_from_slice(plain); + return Ok(()); + } + let mut pos = 0usize; + for &valid in def_levels { + if !valid { + match layout { + ValueLayout::Fixed(width) => out.extend(std::iter::repeat_n(0u8, width)), + ValueLayout::LengthPrefixed => out.extend_from_slice(&[0u8; 4]), + } + continue; + } + let width = match layout { + ValueLayout::Fixed(width) => width, + ValueLayout::LengthPrefixed => { + let len = plain + .get(pos..pos + 4) + .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as usize) + .ok_or_else(|| { + IoError::Format("value section shorter than its definition levels".into()) + })?; + 4 + len + } + }; + let entry = plain.get(pos..pos + width).ok_or_else(|| { + IoError::Format("value section shorter than its definition levels".into()) + })?; + out.extend_from_slice(entry); + pos += width; + } + Ok(()) +} + +/// Split a 4-byte little-endian length off the front of a level or RLE +/// boolean stream and return the run it introduces. +fn length_prefixed_run(buf: &[u8]) -> Result<&[u8], IoError> { + let len = buf + .get(..4) + .map(|b| u32::from_le_bytes(b.try_into().unwrap()) as usize) + .ok_or_else(|| IoError::Format("truncated length-prefixed run".into()))?; + buf.get(4..4 + len) + .ok_or_else(|| IoError::Format("truncated length-prefixed run".into())) +} + +/// DataPageV1 reader. +/// +/// The whole page body is compressed as one unit. Inside it, repetition +/// levels (absent for a flat schema) and definition levels (absent for a +/// REQUIRED column) each lead with a 4-byte length, and the values follow. +/// Returns the definition levels, the value encoding and the raw values. +fn read_data_page_v1( + r: &mut R, + ph: &PageHeader, + codec: Option, + max_def_level: u8, +) -> Result<(Vec, ParquetEncoding, Vec), IoError> { + let h = ph + .data_page_header + .as_ref() + .ok_or_else(|| IoError::Format("missing DataPageHeader".into()))?; + + let mut body = vec![0u8; ph.compressed_page_size as usize]; + r.read_exact(&mut body)?; + let body = match codec { + Some(c) => decompress(&body, c)?, + None => body, + }; + + let n = h.num_values as usize; + let mut pos = 0usize; + let def_levels = if max_def_level == 0 { + vec![true; n] + } else { + if h.definition_level_encoding != ParquetEncoding::Rle { + return Err(IoError::UnsupportedEncoding(format!( + "{:?} definition levels", + h.definition_level_encoding + ))); + } + let run = length_prefixed_run(&body[pos..])?; + pos += 4 + run.len(); + decode_hybrid(run, max_def_level, n)? + .into_iter() + .map(|v| v == max_def_level as u32) + .collect() + }; + + Ok((def_levels, h.encoding, body[pos..].to_vec())) +} + +/// DataPageV2 reader. +/// +/// Repetition and definition levels sit uncompressed at the front of the +/// page with their byte lengths in the header. The value section that +/// follows is compressed only when the header says so. Returns the +/// definition levels, the value encoding and the raw values. fn read_data_page_v2( r: &mut R, ph: &PageHeader, - cmeta: &ColumnMetadata, + codec: Option, + max_def_level: u8, ) -> Result<(Vec, ParquetEncoding, Vec), IoError> { let h = ph .data_page_header_v2 .as_ref() .ok_or_else(|| IoError::Format("missing DataPageHeaderV2".into()))?; - // 1) consume repetition‐levels bytes + // 1) consume repetition-levels bytes. A flat schema has no repetition + // levels, though a writer may still emit a run for them. let mut rep = vec![0u8; h.repetition_levels_byte_length as usize]; r.read_exact(&mut rep)?; - // 2) consume definition‐levels bytes + // 2) consume definition-levels bytes let mut def = vec![0u8; h.definition_levels_byte_length as usize]; r.read_exact(&mut def)?; @@ -287,54 +608,35 @@ fn read_data_page_v2( r.read_exact(&mut vs)?; // 4) decompress if needed - let values_raw = match (h.is_compressed, map_codec(cmeta.codec)) { + let values_raw = match (h.is_compressed, codec) { (true, Some(c)) => decompress(&vs, c)?, _ => vs, }; - // 5) decode definition‐levels (we ignore repetition entirely) - let def_levels = if cmeta.definition_level == 0 && def.is_empty() { - vec![true; h.num_rows as usize] + // 5) decode definition levels. num_values counts every row of the + // page, nulls included, which is the definition level count. + let n = h.num_values as usize; + let def_levels = if max_def_level == 0 || def.is_empty() { + vec![true; n] } else { - decode_hybrid(&def, 1, h.num_rows as usize)? + decode_hybrid(&def, max_def_level, n)? .into_iter() - .map(|v| v != 0) + .map(|v| v == max_def_level as u32) .collect() }; Ok((def_levels, h.encoding, values_raw)) } -/// DataPageV1 reader -fn read_data_page_v1( - r: &mut R, - _ph: &PageHeader, - cmeta: &ColumnMetadata, -) -> Result<(Vec, ParquetEncoding, Vec), IoError> { - // read the 4-byte prefix of the def‐levels stream - let def = read_len_prefixed(r)?; - // read the remaining compressed values for this page, - // which in V1 is “to the end of this page” - as V1 only writes one page - let mut vs = Vec::new(); - r.read_to_end(&mut vs)?; - - let num_vals = (cmeta.num_values as usize).max(def_levels_count(&def, 1)); - let def_levels = if cmeta.definition_level == 0 && def.is_empty() { - vec![true; num_vals] - } else { - decode_hybrid(&def, 1, num_vals)? - .into_iter() - .map(|v| v != 0) - .collect() - }; - Ok((def_levels, ParquetEncoding::Plain, vs)) -} - // Column-value decoder +/// Build the column array from PLAIN entries, one per row. +/// +/// `buf` holds one entry per row in the layout the PLAIN decoders expect, +/// with a zero entry under each null. Categorical columns pass their +/// merged dictionary in `dict` and hold u32 indices in `buf`. fn decode_column( ty: &ArrowType, - enc: ParquetEncoding, dict: &[Vec], buf: &[u8], len: usize, @@ -344,50 +646,33 @@ fn decode_column( Ok(match ty { // numerics - ArrowType::Int32 if enc == ParquetEncoding::Plain => { - Array::NumericArray(NumericArray::Int32(Arc::new(IntegerArray::from_vec64( - decode_int32_plain(buf)?, - mask, - )))) - } - ArrowType::UInt32 if enc == ParquetEncoding::Plain => { - Array::NumericArray(NumericArray::UInt32(Arc::new(IntegerArray::from_vec64( - decode_uint32_as_int32_plain(buf)?, - mask, - )))) - } - ArrowType::Int64 if enc == ParquetEncoding::Plain => { - Array::NumericArray(NumericArray::Int64(Arc::new(IntegerArray::from_vec64( - decode_int64_plain(buf)?, - mask, - )))) - } - ArrowType::UInt64 if enc == ParquetEncoding::Plain => { - Array::NumericArray(NumericArray::UInt64(Arc::new(IntegerArray::from_vec64( - decode_uint64_as_int64_plain(buf)?, - mask, - )))) - } - ArrowType::Float32 if enc == ParquetEncoding::Plain => { - Array::NumericArray(NumericArray::Float32(Arc::new(FloatArray::from_vec64( - decode_float32_plain(buf)?, - mask, - )))) - } - ArrowType::Float64 if enc == ParquetEncoding::Plain => { - Array::NumericArray(NumericArray::Float64(Arc::new(FloatArray::from_vec64( - decode_float64_plain(buf)?, - mask, - )))) - } - - // booleans - ArrowType::Boolean if enc == ParquetEncoding::Plain => { - Array::BooleanArray(Arc::new(BooleanArray::new(Bitmask::from_bytes(buf, len), mask))) + ArrowType::Int32 => Array::NumericArray(NumericArray::Int32(Arc::new( + IntegerArray::from_vec64(decode_int32_plain(buf)?, mask), + ))), + ArrowType::UInt32 => Array::NumericArray(NumericArray::UInt32(Arc::new( + IntegerArray::from_vec64(decode_uint32_as_int32_plain(buf)?, mask), + ))), + ArrowType::Int64 => Array::NumericArray(NumericArray::Int64(Arc::new( + IntegerArray::from_vec64(decode_int64_plain(buf)?, mask), + ))), + ArrowType::UInt64 => Array::NumericArray(NumericArray::UInt64(Arc::new( + IntegerArray::from_vec64(decode_uint64_as_int64_plain(buf)?, mask), + ))), + ArrowType::Float32 => Array::NumericArray(NumericArray::Float32(Arc::new( + FloatArray::from_vec64(decode_float32_plain(buf)?, mask), + ))), + ArrowType::Float64 => Array::NumericArray(NumericArray::Float64(Arc::new( + FloatArray::from_vec64(decode_float64_plain(buf)?, mask), + ))), + + // booleans, one byte per row + ArrowType::Boolean => { + let bits: Vec = buf.iter().map(|&b| b != 0).collect(); + Array::BooleanArray(Arc::new(BooleanArray::new(Bitmask::from_bools(&bits), mask))) } // strings - ArrowType::String if enc == ParquetEncoding::Plain => { + ArrowType::String => { let (offsets, data) = decode_string_plain(buf, len)?; Array::TextArray(TextArray::String32(Arc::new(StringArray { offsets: offsets.into(), @@ -396,7 +681,7 @@ fn decode_column( }))) } #[cfg(feature = "large_string")] - ArrowType::LargeString if enc == ParquetEncoding::Plain => { + ArrowType::LargeString => { use crate::models::decoders::parquet::decode_large_string_plain; let (offsets, data) = decode_large_string_plain(buf, len)?; @@ -407,84 +692,88 @@ fn decode_column( }))) } - // dictionary / categoricals - ArrowType::Dictionary(key_ty) => { - match (key_ty, enc) { - // u32 keys - #[cfg(any( - not(feature = "default_categorical_8"), - feature = "extended_categorical" - ))] - (CategoricalIndexType::UInt32, ParquetEncoding::RleDictionary) => { - let idx = decode_dictionary_indices_rle(buf, len)?; - build_cat32(idx, dict, mask) - } - #[cfg(any( - not(feature = "default_categorical_8"), - feature = "extended_categorical" - ))] - (CategoricalIndexType::UInt32, ParquetEncoding::Plain) => { - let idx = decode_uint32_as_int32_plain(buf)?; - build_cat32(idx, dict, mask) - } - - // u8 keys (default when default_categorical_8 is enabled) - #[cfg(feature = "default_categorical_8")] - (CategoricalIndexType::UInt8, ParquetEncoding::RleDictionary) => { - let idx = decode_dictionary_indices_rle(buf, len)?; - build_cat8(idx, dict, mask) - } - #[cfg(feature = "default_categorical_8")] - (CategoricalIndexType::UInt8, ParquetEncoding::Plain) => { - let idx = decode_uint32_as_int32_plain(buf)?; - build_cat8(idx, dict, mask) - } - - // optional u64 keys - #[cfg(all(feature = "extended_categorical", feature = "large_string"))] - (CategoricalIndexType::UInt64, ParquetEncoding::RleDictionary) => { - let idx = decode_dictionary_indices_rle(buf, len)?; - let idx = idx.into_iter().map(|v| v as u64).collect(); - build_cat64(idx, dict, mask) - } - #[cfg(all(feature = "extended_categorical", feature = "large_string"))] - (CategoricalIndexType::UInt64, ParquetEncoding::Plain) => { - let idx = decode_uint64_as_int64_plain(buf)?; - build_cat64(idx, dict, mask) - } - - _ => { - return Err(IoError::UnsupportedEncoding(format!( - "{:?} + {:?}", - key_ty, enc - ))); - } + // dictionary / categoricals, u32 indices per row + ArrowType::Dictionary(key_ty) => match key_ty { + #[cfg(any( + not(feature = "default_categorical_8"), + feature = "extended_categorical" + ))] + CategoricalIndexType::UInt32 => { + build_cat32(decode_uint32_as_int32_plain(buf)?, dict, mask) } - } + #[cfg(feature = "default_categorical_8")] + CategoricalIndexType::UInt8 => build_cat8(decode_uint32_as_int32_plain(buf)?, dict, mask), + #[cfg(all(feature = "extended_categorical", feature = "large_string"))] + CategoricalIndexType::UInt64 => { + let idx = decode_uint32_as_int32_plain(buf)? + .into_iter() + .map(|v| v as u64) + .collect(); + build_cat64(idx, dict, mask) + } + // Which index widths exist depends on minarrow's categorical + // feature flags, so this arm is unreachable in some builds. + #[allow(unreachable_patterns)] + _ => { + return Err(IoError::UnsupportedType(format!( + "dictionary index {:?}", + key_ty + ))); + } + }, // temporal #[cfg(feature = "datetime")] - ArrowType::Date32 if enc == ParquetEncoding::Plain => { - Array::TemporalArray(TemporalArray::Datetime32(Arc::new(DatetimeArray { + ArrowType::Date32 => Array::TemporalArray(TemporalArray::Datetime32(Arc::new( + DatetimeArray { data: decode_datetime32_plain(buf)?.into(), null_mask: mask, time_unit: Default::default(), - }))) - } + }, + ))), #[cfg(feature = "datetime")] - ArrowType::Date64 if enc == ParquetEncoding::Plain => { - Array::TemporalArray(TemporalArray::Datetime64(Arc::new(DatetimeArray { + ArrowType::Date64 => Array::TemporalArray(TemporalArray::Datetime64(Arc::new( + DatetimeArray { data: decode_datetime64_plain(buf)?.into(), null_mask: mask, time_unit: Default::default(), + }, + ))), + + // decimals + #[cfg(feature = "decimal")] + ArrowType::Decimal32(precision, scale) => { + let data = decode_int32_plain(buf)?; + Array::NumericArray(NumericArray::Decimal32(Arc::new(DecimalArray { + data: data.into(), + null_mask: mask, + precision: *precision, + scale: *scale, + }))) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal64(precision, scale) => { + let data = decode_int64_plain(buf)?; + Array::NumericArray(NumericArray::Decimal64(Arc::new(DecimalArray { + data: data.into(), + null_mask: mask, + precision: *precision, + scale: *scale, + }))) + } + #[cfg(feature = "decimal")] + ArrowType::Decimal128(precision, scale) => { + let data = decode_decimal128_plain(buf)?; + Array::NumericArray(NumericArray::Decimal128(Arc::new(DecimalArray { + data: data.into(), + null_mask: mask, + precision: *precision, + scale: *scale, }))) } _ => { - return Err(IoError::UnsupportedType(format!( - "decode {:?} / {:?}", - ty, enc - ))); + return Err(IoError::UnsupportedType(format!("decode {:?}", ty))); } }) } @@ -607,41 +896,29 @@ fn read_uleb128(buf: &[u8]) -> Result<(u64, usize), IoError> { Err(IoError::Format("ULEB128 overflow/truncate".into())) } -// utility for legacy V1 count heuristic -fn def_levels_count(buf: &[u8], bw: u8) -> usize { - if buf.is_empty() { - 0 - } else if buf[0] & 1 == 0 { - ((buf[0] as usize) >> 1).min(1 << bw) - } else { - 0 - } -} // Misc helpers -fn map_codec(id: i32) -> Option { +/// Resolve the column chunk's codec id. Uncompressed is `None`. Codecs +/// outside the enabled compression features are an error, since their +/// pages cannot be decoded. +fn map_codec(id: i32) -> Result, IoError> { match id { - 0 => None, + 0 => Ok(None), #[cfg(feature = "snappy")] - 1 => Some(Compression::Snappy), + 1 => Ok(Some(Compression::Snappy)), #[cfg(feature = "zstd")] - 6 => Some(Compression::Zstd), // spec: ZSTD = 6 - _ => None, + 6 => Ok(Some(Compression::Zstd)), // spec: ZSTD = 6 + other => Err(IoError::Compression(format!( + "unsupported parquet compression codec {other}" + ))), } } -fn read_len_prefixed(r: &mut R) -> Result, IoError> { - let mut l4 = [0u8; 4]; - r.read_exact(&mut l4)?; - let len = u32::from_le_bytes(l4) as usize; - let mut buf = vec![0u8; len]; - r.read_exact(&mut buf)?; - Ok(buf) -} +/// Split a PLAIN BYTE_ARRAY dictionary page body into its entries. fn parse_dictionary_values(buf: &[u8]) -> Result>, IoError> { - let mut c = std::io::Cursor::new(buf); + let mut c = Cursor::new(buf); let mut out = Vec::new(); while (c.position() as usize) < buf.len() { let mut l4 = [0u8; 4]; @@ -889,7 +1166,6 @@ fn parse_column_meta_data(r: &mut R) -> Result data_page_offset: data_page_offset.unwrap_or(0), dictionary_page_offset, statistics, - definition_level: 0, }) } @@ -1314,14 +1590,12 @@ mod tests { fn decode_column_categorical_rle_dictionary() { let dict_raw = dict(&["foo", "bar"]); let idx: Vec = vec![0, 1, 1, 0]; - let mut encoded = Vec::new(); - encode_dictionary_indices_rle(&idx, &mut encoded).unwrap(); + let encoded: Vec = idx.iter().flat_map(|v| v.to_le_bytes()).collect(); let def_levels = vec![true; idx.len()]; let array = super::decode_column( &ArrowType::Dictionary(CategoricalIndexType::UInt32), - ParquetEncoding::RleDictionary, &dict_raw, &encoded, idx.len(), @@ -1344,14 +1618,12 @@ mod tests { fn decode_column_categorical_rle_dictionary() { let dict_raw = dict(&["foo", "bar"]); let idx: Vec = vec![0, 1, 1, 0]; - let mut encoded = Vec::new(); - encode_dictionary_indices_rle(&idx, &mut encoded).unwrap(); + let encoded: Vec = idx.iter().flat_map(|v| v.to_le_bytes()).collect(); let def_levels = vec![true; idx.len()]; let array = super::decode_column( &ArrowType::Dictionary(CategoricalIndexType::UInt8), - ParquetEncoding::RleDictionary, &dict_raw, &encoded, idx.len(), @@ -1379,14 +1651,7 @@ mod tests { } let def_levels = vec![true; values.len()]; - let array = decode_column( - &ArrowType::Int32, - ParquetEncoding::Plain, - &[], - &buf, - values.len(), - def_levels.clone(), - ) + let array = decode_column(&ArrowType::Int32, &[], &buf, values.len(), def_levels.clone()) .unwrap(); match array { @@ -1401,16 +1666,10 @@ mod tests { #[test] fn decode_column_boolean_plain() { let bits = [true, false, true, true, false, false]; - let data_mask = Bitmask::from_bools(&bits); - let def_levels = bits.to_vec(); // no nulls - let array = decode_column( - &ArrowType::Boolean, - ParquetEncoding::Plain, - &[], - data_mask.as_slice(), - bits.len(), - def_levels, - ) + // one byte per row, as the page readers unpack boolean pages + let bytes: Vec = bits.iter().map(|&b| b as u8).collect(); + let def_levels = vec![true; bits.len()]; + let array = decode_column(&ArrowType::Boolean, &[], &bytes, bits.len(), def_levels) .unwrap(); match array { diff --git a/rust/src/models/types/parquet.rs b/rust/src/models/types/parquet.rs index b95782a..ea49e5c 100644 --- a/rust/src/models/types/parquet.rs +++ b/rust/src/models/types/parquet.rs @@ -33,6 +33,8 @@ pub(crate) enum ParquetPhysicalType { Double = 5, /// Variable-length byte array (used for strings and binary data). ByteArray = 6, + /// Fixed-length byte array (used for Decimal128 and other fixed-width types). + FixedLenByteArray = 7, } impl ParquetPhysicalType { @@ -52,6 +54,7 @@ impl ParquetPhysicalType { 4 => Some(Self::Float), 5 => Some(Self::Double), 6 => Some(Self::ByteArray), + 7 => Some(Self::FixedLenByteArray), _ => None, } } @@ -98,6 +101,14 @@ pub(crate) enum ParquetLogicalType { /// Whether the type is signed (`true`) or unsigned (`false`). is_signed: bool, }, + /// Fixed-point decimal with precision and scale. + #[cfg(feature = "decimal")] + DecimalType { + /// Total number of significant digits. + precision: u8, + /// Digits after the decimal point. + scale: i8, + }, } impl ParquetLogicalType { @@ -363,11 +374,20 @@ pub(crate) fn arrow_type_to_parquet( #[cfg(feature = "datetime")] ArrowType::Interval(_) => panic!("Interval does not map to a parquet type."), #[cfg(feature = "decimal")] - ArrowType::Decimal32(_, _) - | ArrowType::Decimal64(_, _) - | ArrowType::Decimal128(_, _) => { - Err(IoError::UnsupportedType(format!("{ty:?}"))) - } + ArrowType::Decimal32(p, s) => Ok(( + ParquetPhysicalType::Int32, + ParquetLogicalType::DecimalType { precision: *p, scale: *s }, + )), + #[cfg(feature = "decimal")] + ArrowType::Decimal64(p, s) => Ok(( + ParquetPhysicalType::Int64, + ParquetLogicalType::DecimalType { precision: *p, scale: *s }, + )), + #[cfg(feature = "decimal")] + ArrowType::Decimal128(p, s) => Ok(( + ParquetPhysicalType::FixedLenByteArray, + ParquetLogicalType::DecimalType { precision: *p, scale: *s }, + )), #[cfg(all(feature = "extended_categorical", feature = "extended_numeric_types"))] &minarrow::ArrowType::Dictionary( minarrow::ffi::arrow_dtype::CategoricalIndexType::UInt16, @@ -495,6 +515,23 @@ pub(crate) fn parquet_to_arrow_type( Ok(ArrowType::Time64(TimeUnit::Nanoseconds)) } + // Decimals + #[cfg(feature = "decimal")] + ( + ParquetPhysicalType::Int32, + Some(ParquetLogicalType::DecimalType { precision, scale }), + ) => Ok(ArrowType::Decimal32(precision, scale)), + #[cfg(feature = "decimal")] + ( + ParquetPhysicalType::Int64, + Some(ParquetLogicalType::DecimalType { precision, scale }), + ) => Ok(ArrowType::Decimal64(precision, scale)), + #[cfg(feature = "decimal")] + ( + ParquetPhysicalType::FixedLenByteArray, + Some(ParquetLogicalType::DecimalType { precision, scale }), + ) => Ok(ArrowType::Decimal128(precision, scale)), + // Floats (ParquetPhysicalType::Float, _) => Ok(ArrowType::Float32), (ParquetPhysicalType::Double, _) => Ok(ArrowType::Float64), diff --git a/rust/src/models/writers/parquet.rs b/rust/src/models/writers/parquet.rs index b12eb06..5567780 100644 --- a/rust/src/models/writers/parquet.rs +++ b/rust/src/models/writers/parquet.rs @@ -49,7 +49,7 @@ use minarrow::TemporalArray; use minarrow::{Array, NumericArray, Table, TextArray}; use crate::compression::{Compression, compress}; -use crate::constants::PARQUET_MAGIC; +use crate::constants::{PARQUET_CREATED_BY, PARQUET_MAGIC}; use crate::error::IoError; #[cfg(feature = "large_string")] use crate::models::encoders::parquet::data::encode_large_string_plain; @@ -63,7 +63,7 @@ use crate::models::encoders::parquet::metadata::{ PageHeader, PageType, RowGroupMeta, SchemaElement, Statistics, }; use crate::models::types::parquet::ParquetLogicalType::{self}; -use crate::models::types::parquet::{ParquetEncoding, arrow_type_to_parquet}; +use crate::models::types::parquet::{ParquetEncoding, ParquetPhysicalType, arrow_type_to_parquet}; // Chunk size for page splitting pub const PARQUET_PAGE_CHUNK_SIZE: usize = 32_768; @@ -113,14 +113,15 @@ pub fn write_parquet_table( }); for (i, c) in table.cols.iter().enumerate() { let (physical, logical) = arrow_type_to_parquet(&c.field.dtype).unwrap(); + let (precision, scale, type_length) = decimal_schema_fields(&logical, physical); schema.push(SchemaElement { name: c.field.name.clone(), repetition_type: if c.field.nullable { 1 } else { 0 }, // OPTIONAL / REQUIRED type_: Some(physical), converted_type: logical_to_converted(&logical), - type_length: None, - precision: None, - scale: None, + type_length, + precision, + scale, field_id: Some(i as i32), num_children: None, }); @@ -193,27 +194,76 @@ pub fn write_parquet_table( let end = usize::min(start + PARQUET_PAGE_CHUNK_SIZE, n); let len = end - start; - // encode the raw values for this slice + // rep / def levels for this chunk + let def_levels = col.array.null_mask().map_or_else( + || vec![true; len], + |mask| (start..end).map(|i| mask.get(i)).collect(), + ); + let has_nulls = def_levels.iter().any(|&v| !v); let mut values_raw = Vec::new(); + + // Encode a fixed-width slice for this page. The Parquet value + // section holds non-null values only, so null slots are dropped + // before encoding. + macro_rules! encode_valid { + ($encode:ident, $data:expr) => { + if has_nulls { + let valid: Vec<_> = $data + .iter() + .zip(&def_levels) + .filter(|(_, valid)| **valid) + .map(|(v, _)| *v) + .collect(); + $encode(&valid, &mut values_raw) + } else { + $encode($data, &mut values_raw) + } + }; + } + + // encode the raw values for this slice match &col.array { Array::NumericArray(n) => match n { NumericArray::Int32(a) => { - encode_int32_plain(&a.data[start..end], &mut values_raw) + encode_valid!(encode_int32_plain, &a.data[start..end]) } NumericArray::UInt32(a) => { - encode_uint32_as_int32_plain(&a.data[start..end], &mut values_raw) + encode_valid!(encode_uint32_as_int32_plain, &a.data[start..end]) } NumericArray::Int64(a) => { - encode_int64_plain(&a.data[start..end], &mut values_raw) + encode_valid!(encode_int64_plain, &a.data[start..end]) } NumericArray::UInt64(a) => { - encode_uint64_as_int64_plain(&a.data[start..end], &mut values_raw) + encode_valid!(encode_uint64_as_int64_plain, &a.data[start..end]) } NumericArray::Float32(a) => { - encode_float32_plain(&a.data[start..end], &mut values_raw) + encode_valid!(encode_float32_plain, &a.data[start..end]) } NumericArray::Float64(a) => { - encode_float64_plain(&a.data[start..end], &mut values_raw) + encode_valid!(encode_float64_plain, &a.data[start..end]) + } + #[cfg(feature = "decimal")] + NumericArray::Decimal32(a) => { + encode_valid!(encode_int32_plain, &a.data[start..end]) + } + #[cfg(feature = "decimal")] + NumericArray::Decimal64(a) => { + encode_valid!(encode_int64_plain, &a.data[start..end]) + } + #[cfg(feature = "decimal")] + NumericArray::Decimal128(a) => { + let slice = &a.data[start..end]; + if has_nulls { + for (v, valid) in slice.iter().zip(&def_levels) { + if *valid { + values_raw.extend_from_slice(&v.to_be_bytes()); + } + } + } else { + for v in slice { + values_raw.extend_from_slice(&v.to_be_bytes()); + } + } } _ => return Err(IoError::UnsupportedType("numeric".into())), }, @@ -253,31 +303,38 @@ pub fn write_parquet_table( Array::TemporalArray(TemporalArray::Datetime32(a)) => { use crate::models::encoders::parquet::data::encode_datetime32_plain; - encode_datetime32_plain(&a.data[start..end], &mut values_raw) + encode_valid!(encode_datetime32_plain, &a.data[start..end]) } #[cfg(feature = "datetime")] Array::TemporalArray(TemporalArray::Datetime64(a)) => { use crate::models::encoders::parquet::data::encode_datetime64_plain; - encode_datetime64_plain(&a.data[start..end], &mut values_raw) + encode_valid!(encode_datetime64_plain, &a.data[start..end]) } #[cfg(any( not(feature = "default_categorical_8"), feature = "extended_categorical" ))] Array::TextArray(TextArray::Categorical32(a)) => { - encode_dictionary_indices_rle(&a.data[start..end], &mut values_raw)? + encode_valid!(encode_dictionary_indices_rle, &a.data[start..end])? } #[cfg(feature = "default_categorical_8")] Array::TextArray(TextArray::Categorical8(a)) => { - let idx: Vec = a.data[start..end].iter().map(|&v| v as u32).collect(); + let idx: Vec = a.data[start..end] + .iter() + .zip(&def_levels) + .filter(|(_, valid)| **valid) + .map(|(&v, _)| v as u32) + .collect(); encode_dictionary_indices_rle(&idx, &mut values_raw)? } #[cfg(all(feature = "extended_categorical", feature = "large_string"))] Array::TextArray(TextArray::Categorical64(a)) => { let idx: Vec = a.data[start..end] .iter() - .map(|&v| u32::try_from(v)) + .zip(&def_levels) + .filter(|(_, valid)| **valid) + .map(|(&v, _)| u32::try_from(v)) .collect::>() .map_err(|_| { IoError::Format( @@ -289,11 +346,6 @@ pub fn write_parquet_table( _ => return Err(IoError::UnsupportedType(format!("array {:?}", col.array))), } - // rep / def levels for this chunk - let def_levels = col.array.null_mask().map_or_else( - || vec![true; len], - |mask| (start..end).map(|i| mask.get(i)).collect(), - ); let def_buf = encode_levels_rle(&def_levels); let rep_buf = encode_levels_rle(&vec![false; len]); @@ -337,7 +389,8 @@ pub fn write_parquet_table( data_page_header_v2: Some(DataPageHeaderV2 { num_rows: len as i32, num_nulls: def_levels.iter().filter(|&&v| !v).count() as i32, - num_values: (len - def_levels.iter().filter(|&&v| !v).count()) as i32, + // num_values counts every row of the page, nulls included. + num_values: len as i32, encoding: if is_dictionary(&col.array) { ParquetEncoding::RleDictionary } else { @@ -389,7 +442,6 @@ pub fn write_parquet_table( data_page_offset: first_data, dictionary_page_offset, statistics: None, - definition_level: if col.field.nullable { 1 } else { 0 }, }, }); } @@ -408,7 +460,7 @@ pub fn write_parquet_table( num_rows: n_rows_i64, row_groups, key_value_metadata: None, - created_by: Some("parquet_writer-v2".into()), + created_by: Some(PARQUET_CREATED_BY.into()), } .write(&mut out)?; @@ -417,6 +469,25 @@ pub fn write_parquet_table( // Helpers +/// Extract precision, scale, and type_length for decimal schema elements. +/// Non-decimal types return `(None, None, None)`. +fn decimal_schema_fields( + logical: &ParquetLogicalType, + #[allow(unused_variables)] physical: ParquetPhysicalType, +) -> (Option, Option, Option) { + #[cfg(feature = "decimal")] + if let ParquetLogicalType::DecimalType { precision, scale } = logical { + let type_length = if matches!(physical, ParquetPhysicalType::FixedLenByteArray) { + Some(16) + } else { + None + }; + return (Some(*precision as i32), Some(*scale as i32), type_length); + } + let _ = logical; + (None, None, None) +} + /// Add a dictionary page and return its uncompressed and compressed byte /// contributions to the column chunk totals, each including the page header. fn write_dictionary_page<'a, W, I>( @@ -602,6 +673,8 @@ fn logical_to_converted(log: &ParquetLogicalType) -> Option { bit_width: 64, is_signed: true, } => 18, + #[cfg(feature = "decimal")] + ParquetLogicalType::DecimalType { .. } => 5, _ => return None, }) } From a8e4fa42ef39334932267ea30317106d769c96dc Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Thu, 10 Sep 2026 05:26:37 +0100 Subject: [PATCH 3/7] Bump to 0.6.2 on minarrow 0.18.1 and make the Parquet reader and writer conformant with pyarrow (TSK500) --- CHANGELOG.md | 19 ++ python/Cargo.lock | 8 +- python/Cargo.toml | 4 +- python/pyproject.toml | 4 +- python/tests/test_files.py | 34 ++ rust/Cargo.lock | 8 +- rust/Cargo.toml | 2 +- rust/README.md | 2 +- .../generate_pyarrow_parquet_files.py | 85 +++++ .../pyarrow_dictionary_v2.parquet | Bin 0 -> 1802 bytes .../pyarrow_nullable_row_groups.parquet | Bin 0 -> 3656 bytes .../pyarrow_plain_v2.parquet | Bin 0 -> 1679 bytes rust/pyarrow-roundtrip/pyarrow_simple.parquet | Bin 0 -> 1105 bytes rust/src/constants.rs | 6 + rust/src/models/encoders/parquet/data.rs | 84 +++-- rust/src/models/encoders/parquet/metadata.rs | 2 - rust/src/models/readers/ipc/window.rs | 8 +- rust/tests/parquet_nullable_roundtrip.rs | 323 ++++++++++++++++++ rust/tests/pyarrow_parquet.rs | 219 ++++++++++++ 19 files changed, 757 insertions(+), 51 deletions(-) create mode 100644 rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py create mode 100644 rust/pyarrow-roundtrip/pyarrow_dictionary_v2.parquet create mode 100644 rust/pyarrow-roundtrip/pyarrow_nullable_row_groups.parquet create mode 100644 rust/pyarrow-roundtrip/pyarrow_plain_v2.parquet create mode 100644 rust/pyarrow-roundtrip/pyarrow_simple.parquet create mode 100644 rust/tests/parquet_nullable_roundtrip.rs create mode 100644 rust/tests/pyarrow_parquet.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 54abc85..235d1ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,25 @@ Notable changes are recorded from 0.5.0 onward. +## 0.6.2 + +### Changed + +- minarrow 0.18.1 and vec64 0.5.1. +- New `decimal` feature forwarding minarrow's `decimal` feature. Decimal32, Decimal64 and Decimal128 columns are supported in Arrow IPC and in Parquet, where they map to the DECIMAL logical type over INT32, INT64 and FIXED_LEN_BYTE_ARRAY. +- The Python package pins minarrow and minarrow-pyo3 at 0.18.1. +- The Parquet writer follows the Parquet value layout for nullable columns. Value sections hold non-null values only and DataPageV2 headers count every row in `num_values`, so files with nulls now read in pyarrow and other Parquet readers. Files with nulls written by earlier releases do not read back under this release. + +### Fixed + +- The Parquet reader failed with `UnexpectedEof` on files written by pyarrow. It now reads DataPageV1 pages with compressed levels, `RLE` booleans, dictionary-encoded columns of any physical type, and files with several row groups. Dictionary pages in files without lightstream's `created_by` marker expand to the schema type rather than being read as categorical columns. + +## 0.6.1 + +### Changed + +- minarrow 0.17.0, vec64 0.5.0, arrow 59.2.0 and polars 0.55.2. + ## 0.6.0 ### Changed diff --git a/python/Cargo.lock b/python/Cargo.lock index dc13912..e9cc009 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -662,7 +662,7 @@ checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "lightstream" -version = "0.6.1" +version = "0.6.2" dependencies = [ "bytes", "fast-float2", @@ -692,7 +692,7 @@ dependencies = [ [[package]] name = "lightstream-py" -version = "0.6.1" +version = "0.6.2" dependencies = [ "futures-core", "lightstream", @@ -745,9 +745,9 @@ dependencies = [ [[package]] name = "minarrow-pyo3" -version = "0.18.0" +version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6363f738664b325aab048cee8101e79606a487d9a8a294b960a1f55220961c4" +checksum = "94e91845f51b0aab761625144ab1c71b989d5237d8cc6159245db07b4b78bc26" dependencies = [ "minarrow", "pyo3", diff --git a/python/Cargo.toml b/python/Cargo.toml index 1463769..4d36ff2 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -2,7 +2,7 @@ cargo-features = ["trim-paths"] [package] name = "lightstream-py" -version = "0.6.1" +version = "0.6.2" edition = "2024" authors = ["Peter G. Bower"] license = "MPL-2.0" @@ -19,7 +19,7 @@ name = "lightstream_py" crate-type = ["cdylib", "rlib"] [dependencies] -lightstream = { version = "0.6", path = "../rust", features = ["csv", "datetime", "extended_categorical", "extended_numeric_types", "json", "mmap", "http", "parquet", "protocol", "quic", "snappy", "stdio", "tcp", "tls", "uds", "webtransport", "websocket", "zstd"] } +lightstream = { version = "0.6", path = "../rust", features = ["csv", "datetime", "decimal", "extended_categorical", "extended_numeric_types", "json", "mmap", "http", "parquet", "protocol", "quic", "snappy", "stdio", "tcp", "tls", "uds", "webtransport", "websocket", "zstd"] } # The categorical and numeric feature set mirrors the minarrow-py build. # minarrow-pyo3's dictionary-index conversion needs the extended features, # and they flow through lightstream's flags so its match arms gate in step diff --git a/python/pyproject.toml b/python/pyproject.toml index 400e7e8..fc5c59a 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "lightstream-io" -version = "0.6.1" +version = "0.6.2" description = "Streaming Arrow I/O for Python - files, sockets, and network transports with zero-copy minarrow interop." readme = "README.md" requires-python = ">=3.9" @@ -20,7 +20,7 @@ classifiers = [ "Topic :: Scientific/Engineering", "Topic :: Software Development :: Libraries", ] -dependencies = ["minarrow>=0.18"] +dependencies = ["minarrow>=0.18.1"] [project.urls] Homepage = "https://github.com/SpaceCell/lightstream" diff --git a/python/tests/test_files.py b/python/tests/test_files.py index 2f7be22..af8d9ae 100644 --- a/python/tests/test_files.py +++ b/python/tests/test_files.py @@ -15,6 +15,7 @@ import lightstream as ls import minarrow import pyarrow as pa +import pyarrow.parquet as pq import pytest @@ -224,6 +225,39 @@ def test_parquet_multi_write_consolidates(tmp_path): assert result.num_rows == 6 +def nullable_table(): + return pa.table( + { + "id": pa.array([1, None, 3, 4, None], type=pa.int64()), + "small": pa.array([None, -2, 3, -4, 5], type=pa.int32()), + "count": pa.array([1, 2, None, 4, 5], type=pa.uint32()), + "name": pa.array(["a", "b", None, "a", "c"], type=pa.string()), + "score": pa.array([1.5, None, 3.5, 4.5, 5.5], type=pa.float64()), + "ratio": pa.array([0.5, 1.5, 2.5, None, 4.5], type=pa.float32()), + "flag": pa.array([True, False, None, True, False], type=pa.bool_()), + "day": pa.array([1, None, 3, 4, 5], type=pa.date32()), + } + ) + + +def test_parquet_reads_pyarrow_file_with_nulls(tmp_path): + path = str(tmp_path / "pyarrow.parquet") + original = nullable_table() + pq.write_table(original, path) + + result = pa.table(ls.read(path).read_all()) + assert result.to_pydict() == original.to_pydict() + + +def test_parquet_written_with_nulls_reads_in_pyarrow(tmp_path): + path = str(tmp_path / "lightstream.parquet") + original = nullable_table() + with ls.write(path) as w: + w.write(original) + + assert pq.read_table(path).to_pydict() == original.to_pydict() + + @pytest.mark.parametrize("codec", ["zstd", "snappy"]) def test_parquet_compression(tmp_path, codec): path = str(tmp_path / f"quotes_{codec}.parquet") diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 3998e04..9238e94 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1924,7 +1924,7 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "lightstream" -version = "0.6.1" +version = "0.6.2" dependencies = [ "arrow", "arrow-flight", @@ -2048,7 +2048,7 @@ checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" [[package]] name = "minarrow" -version = "0.17.0" +version = "0.18.1" dependencies = [ "arrow", "arrow-schema", @@ -4410,9 +4410,9 @@ dependencies = [ [[package]] name = "vec64" -version = "0.5.0" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94b5a35bf381f20a9d41312bd977f43e997ab5bc82deb85b503b3346e6a437b7" +checksum = "1aef6bbef159f21ac387220ebc71141524423f7886afd390bc7b140f12630425" dependencies = [ "libc", ] diff --git a/rust/Cargo.toml b/rust/Cargo.toml index e7a62a3..be37ee2 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lightstream" -version = "0.6.1" +version = "0.6.2" edition = "2024" license = "MPL-2.0" keywords = [ diff --git a/rust/README.md b/rust/README.md index 42c613e..98e9d64 100644 --- a/rust/README.md +++ b/rust/README.md @@ -200,7 +200,7 @@ cargo run --example uds_arrow --features uds | `msgpack`, `protobuf` | Message encodings (enables `protocol`) | | `tls` | TLS for TCP, WebSocket, HTTP/2 | | `io_uring` | Linux async I/O (experimental) | -| `datetime`, `large_string`, `extended_numeric_types` | Schema extensions | +| `datetime`, `decimal`, `large_string`, `extended_numeric_types` | Schema extensions | ## Buffering and safety diff --git a/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py b/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py new file mode 100644 index 0000000..be5530b --- /dev/null +++ b/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Generate the PyArrow Parquet fixture files for tests/pyarrow_parquet.rs. + +Each file exercises one combination of the page layouts pyarrow can write, +so the reader is checked against real parquet-cpp output rather than +lightstream's own encoder: + + - pyarrow_simple.parquet: three required columns written with pyarrow's + defaults (Snappy, dictionary encoding, DataPageV1) + - pyarrow_nullable_row_groups.parquet: nullable columns with nulls, split + across three row groups, written with pyarrow's defaults + - pyarrow_plain_v2.parquet: the same nullable data with dictionary + encoding off, DataPageV2 and no compression + - pyarrow_dictionary_v2.parquet: the same nullable data with dictionary + encoding on, DataPageV2 and Snappy + +The fixtures are committed alongside this script. Run from rust/ to +regenerate them: + + python3 pyarrow-roundtrip/generate_pyarrow_parquet_files.py +""" + +import os + +import pyarrow as pa +import pyarrow.parquet as pq + +OUT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def simple_table() -> pa.Table: + """The three-column table from the original defect report.""" + return pa.table( + { + "id": pa.array([1, 2, 3, 4, 5], type=pa.int64()), + "name": pa.array(["Alice", "Bob", "Charlie", "Diana", "Eve"], type=pa.string()), + "score": pa.array([85.5, 92.0, 78.3, 95.1, 88.7], type=pa.float64()), + } + ) + + +def nullable_table() -> pa.Table: + """Seven rows over the basic types, matching expected_nullable_table in Rust.""" + return pa.table( + { + "int32": pa.array([1, None, 3, 4, None, 6, 7], type=pa.int32()), + "int64": pa.array([100, 101, None, 103, 104, 105, None], type=pa.int64()), + "float32": pa.array([0.5, 1.5, 2.5, None, 4.5, 5.5, 6.5], type=pa.float32()), + "float64": pa.array([None, -1.0, -2.0, -3.0, -4.0, None, -6.0], type=pa.float64()), + "bool": pa.array([True, False, None, True, False, True, None], type=pa.bool_()), + "string": pa.array(["a", "b", "a", None, "c", "b", "a"], type=pa.string()), + } + ) + + +def main() -> None: + simple = simple_table() + pq.write_table(simple, os.path.join(OUT_DIR, "pyarrow_simple.parquet")) + + nullable = nullable_table() + pq.write_table( + nullable, + os.path.join(OUT_DIR, "pyarrow_nullable_row_groups.parquet"), + row_group_size=3, + ) + pq.write_table( + nullable, + os.path.join(OUT_DIR, "pyarrow_plain_v2.parquet"), + use_dictionary=False, + data_page_version="2.0", + compression="none", + ) + pq.write_table( + nullable, + os.path.join(OUT_DIR, "pyarrow_dictionary_v2.parquet"), + data_page_version="2.0", + compression="snappy", + ) + for name in sorted(os.listdir(OUT_DIR)): + if name.startswith("pyarrow_") and name.endswith(".parquet"): + print(f"wrote {name}") + + +if __name__ == "__main__": + main() diff --git a/rust/pyarrow-roundtrip/pyarrow_dictionary_v2.parquet b/rust/pyarrow-roundtrip/pyarrow_dictionary_v2.parquet new file mode 100644 index 0000000000000000000000000000000000000000..0f2fe5024d9d0ea75eca2275be35857bd7450d86 GIT binary patch literal 1802 zcma)7OHUI~6ux&lOh+$FfP+1=JM5F4Tkrvq zTf0E5Nwfp?&Sx3!_`_}*mY!b1!2`}kYeU=o4H^{6` zFLW;qH{pKd?kDc&eEaa;^UhF0gaIg|$za2$EzdXG;hp@NEq;dvwBJ{T`-}VU41b#^ zhUtHZfyHiL?d!4weH&;k&7%D;-+%G=(dPd)!2_0% z|7xrQv6>lUp0K>2WSYVQ_kq7-4a>d#<|UDEYhS{hFx0Tn+Ql30ZgPLq_P@rHEseG> z4%hKeJI&htX}byEcE8)qOG-x-3b&>&2c0Nin&(OSD`}-v9)skjQVo@^sB219)M!QD zrC~5hyHTZX(28@k;#8!oN%HWFb`vzg14KUxT$c%o6zKwKf=2`#$~dj%{=&1;^x(2m znIhR2a+n%W)JR9RmY*!_{92b9~tTMSU@G+p!fRfsaEMATky22drME(i+A&V zQz~;iv>#Oy&#}FTM`IZ{)1>J@+d#@#(o3`BmlozSlbSNvA|J(&HPrb`zKOq}U+u+5 zew6~(-l<2#C!<1prS53^yh`UO-^fcDS_pz~29dv^lc0;xEFDbfcD(u|iVN>if{IEv zf_iDsxAm+iCVN&&^~X=<*Upw#R?e2{^;)B6eZIcBSX=Jt8$8#0?t)jv56qqq`-ZPz I4*!CG0I19*v;Y7A literal 0 HcmV?d00001 diff --git a/rust/pyarrow-roundtrip/pyarrow_nullable_row_groups.parquet b/rust/pyarrow-roundtrip/pyarrow_nullable_row_groups.parquet new file mode 100644 index 0000000000000000000000000000000000000000..472c26d86cc2d7f1bc22c2f87ef1393bbfb3035d GIT binary patch literal 3656 zcmb7HPe@!x7@v9jeEYOpvx+m!!X7;6V8K|UaY=(vXEn94nz|+>F}HPNSBw9uMw5_3 z2&IQoLMT1$Gp67Q7T;A=gvT6)zt-}_Ff12 zR2xk>psJVkS&iB88?b?G*$i>bzTFYNJqs2Rq)|&LV@a_5Ecl1MX9mH*yM)x>*E<3J z_b?IlAEm;@Gn@>-XjfWmM<4p{Q76?$ngcL?A3ZfAsl}l2?FkXgNDTvsOQQjJ7yx0N ztQ(yq24qSFMR@h7Pp5?jHbucJ2*|?P(T7Bov0G~Pr$7eilAz^?HIEv&8mLArJoq)= z7cNSLi~C4*eHGDYsUBuSmH^>aNjn{1fos@S&4B6~negq7@T@sKqn}&rO8H zh#F>$EtwW%a9Twp{=c%Z(2%fGu}-j0kV1 zva-BhW*NY+-w|!A8@r<;08=n8oF&YMdHUnE&9dI<2v@jAVC61tZmxl!F!gqJw^ui+ z{bUxg;GMbg+?phL=H$kRm(2z0nZrH)4dvfbc~1JrSMS`Q;TL($GKY-6<>jPjJ8%EL5pt~WCzg_v*g*Z&696hzOV~Tj9tq?Ji>-R{`aHeCp!mr8u-neM4?zuc5 z4oT63cqZ^QU|OrV%Ri&?L0Z1{@|bNwd5`(Qy&Q0-qvqeIxhH-8Tbi7ZAP;C+O!8(&I@LKE)}FYxp7 z%bsxrps^U($m2{qe4F2qM?LackFEv5mnUYJJGvGi0sOFyJ3t`+0t!=JWYogKhgZ*= z&yu+3zKw0PT2xAv;D_x4rS?^^s@_woswS&iR)hm{IGa@JB95@OBdiutG=~~)<7@^y z&;U(eGO8Dyst&K^p33FcOwH+#AzHv=F^C$rOH-xL{))cUY#;Dd3Su3;^a;i@rGoQHeWcHq zRI!A311=TVg$B-tPLwWc5HtxPb#Mfw(`z@aUvM8NsH)<=;Y(}2tb5%N(YsaN{`BVZ t?&-?b*6H&0_U3-?-tzX&>gGo8$k@>E(3#vC{I${aj~qhjn}`2H{tM%7%OU^( literal 0 HcmV?d00001 diff --git a/rust/pyarrow-roundtrip/pyarrow_plain_v2.parquet b/rust/pyarrow-roundtrip/pyarrow_plain_v2.parquet new file mode 100644 index 0000000000000000000000000000000000000000..cc6d28bd0018438d248bd22949d5e784c2e86fb7 GIT binary patch literal 1679 zcmaJ?-A)=o6rNqzbz_LpCbOAMxZuW(P1Gu&=?_Vn6-%|%x@co*nnqYD(C}mN2W@)M z3m>2tUfOr)8#8C#V67z)l`S znK4EPVTJ~~ZOrtYrl#O(q>hRaSOzQ;7McUn4&50T{iTo@yDk_`E(Vpj@!M`}3VmIyNsgG#AhpbMni{d2)rfX!<2=COU=I2+`l@++nb;} z^K=_fkEnRZ`vGr7d|^PuUHeKXZ_ZbuyGhse+VGiFb?7BJ#OS%l|uZ+4%`nPT6KM@WgX$2s?5 zA-&?f&Al>z!Fidll*OF~1{)}?aK45d&LM|W9xe*pz!k|zL&x{3{}Lzd3ODz7)K*C{y|161%atW zv$f+eE_e?uD08nZ`z4UK#V$6&Vh4rRm*djWTyStOS7^1G?bu1Fby#gSV)3Q literal 0 HcmV?d00001 diff --git a/rust/pyarrow-roundtrip/pyarrow_simple.parquet b/rust/pyarrow-roundtrip/pyarrow_simple.parquet new file mode 100644 index 0000000000000000000000000000000000000000..438815d0c0e9468465f3f634f0554251d1de07bf GIT binary patch literal 1105 zcmb7E&2G~`5T3Oi$2CDpk=C*yhg@m}H4U^S98fA{oHPOv+B7H)aKd#>64CtNl!PPV z#GRKxJOBp-6$c)GN8rGb1GD}?CKXE718x|N?$5Sq;)3CZyP|&$NW*k zyxA8>csB5SnH*<)k|tB@yTRGOXV;4Lfje?VW`E`{c5xdDGKI}lSiX{}#Sapu0sv#$ z(@vG8vjE-dE@OA!J*++x^08a}@|mP~TwNqFwd97W?Fyk9G3s)3RRdPD45m8u`u@=6 zD$TwjnbP1Cc!nHh^^QG5d#P|6;kNB`<*)CF^fEJ%)>h zN!orR*&h`9t4xhAZ|w5UTcEK-0m&S7Fc-n(7iEd_-^O zVCD@wf{*YX3mZwV-bw16b{JZ$-xtx7VdMNs)H$7xB77nt@Ca`n^&KW7WC!~z@n|$e zw=2MUSP$zxF>fqkyRpBw2WEca1}|Rv=Ud)nvgHQBc%Gm6!E`Vl<#)=Zo26~5LO<8& Mf$#LUZqdi&H$Jw^n*aa+ literal 0 HcmV?d00001 diff --git a/rust/src/constants.rs b/rust/src/constants.rs index e44ba5f..e530085 100644 --- a/rust/src/constants.rs +++ b/rust/src/constants.rs @@ -142,3 +142,9 @@ pub const METADATA_SIZE_PREFIX: usize = 4; /// Required “PAR1” marker written at Parquet file head and tail. pub const PARQUET_MAGIC: &[u8; 4] = b"PAR1"; + +/// Writer identity recorded in the `created_by` field of the Parquet footer. +/// +/// The Parquet reader checks this marker before applying lightstream's +/// categorical column convention, which other writers do not share. +pub const PARQUET_CREATED_BY: &str = "parquet_writer-v2"; diff --git a/rust/src/models/encoders/parquet/data.rs b/rust/src/models/encoders/parquet/data.rs index 00268cb..2014fdb 100644 --- a/rust/src/models/encoders/parquet/data.rs +++ b/rust/src/models/encoders/parquet/data.rs @@ -81,20 +81,23 @@ pub fn encode_float64_plain(data: &[f64], out: &mut Vec) { // // We only support RLE encoding for Categorical types at the present time. -/// Encode a boolean column bit-packed (LSB-first), respecting `null_mask`, appending to `out`. +/// Encode a boolean column bit-packed (LSB-first), appending to `out`. +/// +/// Null slots are omitted, so the output holds one bit per non-null value +/// as the Parquet value section requires. pub fn encode_bool_bitpacked( values: &Bitmask, null_mask: Option<&Bitmask>, len: usize, out: &mut Vec, ) { - //out.clear(); let mut byte = 0u8; let mut bit = 0; for i in 0..len { - let valid = null_mask.is_none_or(|m| m.get(i)); - let v = if valid { values.get(i) } else { false }; - if v { + if !null_mask.is_none_or(|m| m.get(i)) { + continue; + } + if values.get(i) { byte |= 1 << bit; } bit += 1; @@ -111,9 +114,10 @@ pub fn encode_bool_bitpacked( // UTF-8 strings -/// Encode String32 (UTF-8) using length-prefix (u32 LE) per row +/// Encode String32 (UTF-8) using length-prefix (u32 LE) per value /// -/// Nulls emit zero length. +/// Null slots are omitted, so the output holds one entry per non-null +/// value as the Parquet value section requires. pub fn encode_string_plain( offsets: &[u32], values: &[u8], @@ -122,23 +126,21 @@ pub fn encode_string_plain( out: &mut Vec, ) -> Result<(), IoError> { for i in 0..len { - // always emit a 4-byte length prefix - let valid = null_mask.is_none_or(|m| m.get(i)); + if !null_mask.is_none_or(|m| m.get(i)) { + continue; + } let start = offsets[i] as usize; let end = offsets[i + 1] as usize; - let s_len = if valid { end - start } else { 0 }; - out.extend_from_slice(&(s_len as u32).to_le_bytes()); - // only write the bytes for non-null - if valid { - out.extend_from_slice(&values[start..end]); - } + out.extend_from_slice(&((end - start) as u32).to_le_bytes()); + out.extend_from_slice(&values[start..end]); } Ok(()) } /// Encode LargeString (i.e., UTF-8, 64-bit offsets) as length-prefix (u32 LE) /// -/// Nulls emit zero length. +/// Null slots are omitted, so the output holds one entry per non-null +/// value as the Parquet value section requires. #[cfg(feature = "large_string")] pub fn encode_large_string_plain( offsets: &[u64], @@ -148,22 +150,20 @@ pub fn encode_large_string_plain( out: &mut Vec, ) -> Result<(), IoError> { for i in 0..len { - let valid = null_mask.is_none_or(|m| m.get(i)); + if !null_mask.is_none_or(|m| m.get(i)) { + continue; + } let start = offsets[i] as usize; let end = offsets[i + 1] as usize; - let s_len = if valid { end - start } else { 0 }; - if valid && s_len > u32::MAX as usize { + let s_len = end - start; + if s_len > u32::MAX as usize { return Err(IoError::InputDataError(format!( "string >4 GiB ({} bytes)", s_len ))); } - // length prefix for every row out.extend_from_slice(&(s_len as u32).to_le_bytes()); - // actual bytes only if non-null - if valid { - out.extend_from_slice(&values[start..end]); - } + out.extend_from_slice(&values[start..end]); } Ok(()) } @@ -388,8 +388,8 @@ mod tests { #[test] fn test_encode_bool_bitpacked_with_nulls() { - // same values but pretend positions 1 and 3 are null ⇒ should be written as false - let values = vec64![true, true, true, false]; + // positions 1 and 3 are null, so only positions 0, 2 and 4 are packed + let values = vec64![true, true, false, false, true]; let mut nulls = Bitmask::new_set_all(values.len(), true); nulls.set_false(1); nulls.set_false(3); @@ -401,12 +401,32 @@ mod tests { &mut buf, ); - // check that bits at 1 and 3 are 0: - let byte = buf[0]; - assert_eq!((byte >> 0) & 1, 1); // idx0 valid, true - assert_eq!((byte >> 1) & 1, 0); // idx1 null -> treated as false - assert_eq!((byte >> 2) & 1, 1); // idx2 - assert_eq!((byte >> 3) & 1, 0); // idx3 null + // three packed bits: idx0 true, idx2 false, idx4 true + assert_eq!(buf, vec![0b101]); + } + + #[test] + fn test_encode_string_plain_omits_nulls() { + let slices = ["foo", "", "rust"]; + let mut offsets = Vec::with_capacity(slices.len() + 1); + offsets.push(0); + let mut values = Vec::new(); + for s in &slices { + values.extend_from_slice(s.as_bytes()); + offsets.push(values.len() as u32); + } + let mut nulls = Bitmask::new_set_all(slices.len(), true); + nulls.set_false(1); + let mut buf = Vec::new(); + encode_string_plain(&offsets, &values, Some(&nulls), slices.len(), &mut buf).unwrap(); + + // two entries: "foo" and "rust", the null slot has no length prefix + let mut expected = Vec::new(); + expected.extend_from_slice(&3u32.to_le_bytes()); + expected.extend_from_slice(b"foo"); + expected.extend_from_slice(&4u32.to_le_bytes()); + expected.extend_from_slice(b"rust"); + assert_eq!(buf, expected); } #[test] diff --git a/rust/src/models/encoders/parquet/metadata.rs b/rust/src/models/encoders/parquet/metadata.rs index e2bb9ae..61075fc 100644 --- a/rust/src/models/encoders/parquet/metadata.rs +++ b/rust/src/models/encoders/parquet/metadata.rs @@ -106,8 +106,6 @@ pub(crate) struct ColumnMetadata { pub dictionary_page_offset: Option, /// Optional per-column statistics. pub statistics: Option, - /// Definition level (REQUIRED/OPTIONAL/REPEATED encoded level). - pub definition_level: i32, } /// Parquet statistics for a column (min/max, null/unique counts). diff --git a/rust/src/models/readers/ipc/window.rs b/rust/src/models/readers/ipc/window.rs index 2dac296..cf6e28b 100644 --- a/rust/src/models/readers/ipc/window.rs +++ b/rust/src/models/readers/ipc/window.rs @@ -20,6 +20,8 @@ use minarrow::{ }; #[cfg(feature = "datetime")] use minarrow::{DatetimeArray, TemporalArray}; +#[cfg(feature = "decimal")] +use minarrow::DecimalArray; /// Build a standalone table for the row window `[offset, offset + len)` /// of a decoded batch. A window covering the whole table returns a @@ -93,7 +95,7 @@ fn window_array(array: &Array, offset: usize, len: usize) -> io::Result { NumericArray::UInt16(arr) => win_int!(UInt16, arr), #[cfg(feature = "decimal")] NumericArray::Decimal32(arr) => { - NumericArray::Decimal32(Arc::new(minarrow::DecimalArray { + NumericArray::Decimal32(Arc::new(DecimalArray { data: window_buffer(&arr.data, offset, len), null_mask: window_mask(arr.null_mask.as_ref(), offset, len), precision: arr.precision, @@ -102,7 +104,7 @@ fn window_array(array: &Array, offset: usize, len: usize) -> io::Result { } #[cfg(feature = "decimal")] NumericArray::Decimal64(arr) => { - NumericArray::Decimal64(Arc::new(minarrow::DecimalArray { + NumericArray::Decimal64(Arc::new(DecimalArray { data: window_buffer(&arr.data, offset, len), null_mask: window_mask(arr.null_mask.as_ref(), offset, len), precision: arr.precision, @@ -111,7 +113,7 @@ fn window_array(array: &Array, offset: usize, len: usize) -> io::Result { } #[cfg(feature = "decimal")] NumericArray::Decimal128(arr) => { - NumericArray::Decimal128(Arc::new(minarrow::DecimalArray { + NumericArray::Decimal128(Arc::new(DecimalArray { data: window_buffer(&arr.data, offset, len), null_mask: window_mask(arr.null_mask.as_ref(), offset, len), precision: arr.precision, diff --git a/rust/tests/parquet_nullable_roundtrip.rs b/rust/tests/parquet_nullable_roundtrip.rs new file mode 100644 index 0000000..16bc8ec --- /dev/null +++ b/rust/tests/parquet_nullable_roundtrip.rs @@ -0,0 +1,323 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! Parquet write-then-read over every supported column type with nulls. +//! +//! Each column spans several data pages with nulls placed on both sides of +//! every page boundary, so the value sections hold non-null entries only +//! and the reader has to scatter them back against the definition levels +//! page by page. + +#[cfg(feature = "parquet")] +mod parquet_nullable_roundtrip_tests { + use std::io::{Cursor, Seek, SeekFrom}; + + use lightstream::compression::Compression; + use lightstream::models::readers::parquet::load_parquet_table; + use lightstream::models::writers::parquet::{PARQUET_PAGE_CHUNK_SIZE, write_parquet_table}; + use minarrow::ffi::arrow_dtype::CategoricalIndexType; + use minarrow::{ + Array, ArrowType, Bitmask, BooleanArray, CategoricalArray, Field, FieldArray, FloatArray, + IntegerArray, MaskedArray, NumericArray, StringArray, Table, TextArray, Vec64, + }; + #[cfg(feature = "datetime")] + use minarrow::{DatetimeArray, TemporalArray}; + + /// Rows per column: two full pages plus a partial third page. + const N_ROWS: usize = 2 * PARQUET_PAGE_CHUNK_SIZE + 13; + + /// Every seventh row from row three is null, plus the rows either side of + /// each page boundary. + fn is_valid(i: usize) -> bool { + let boundary = [ + PARQUET_PAGE_CHUNK_SIZE - 1, + PARQUET_PAGE_CHUNK_SIZE, + 2 * PARQUET_PAGE_CHUNK_SIZE - 1, + 2 * PARQUET_PAGE_CHUNK_SIZE, + ]; + i % 7 != 3 && !boundary.contains(&i) + } + + fn null_mask() -> Bitmask { + let bits: Vec = (0..N_ROWS).map(is_valid).collect(); + Bitmask::from_bools(&bits) + } + + fn expected(value: impl Fn(usize) -> T) -> Vec> { + (0..N_ROWS) + .map(|i| is_valid(i).then(|| value(i))) + .collect() + } + + fn int32_value(i: usize) -> i32 { + i as i32 - 1000 + } + fn uint32_value(i: usize) -> u32 { + i as u32 * 3 + } + fn int64_value(i: usize) -> i64 { + -(i as i64) * 7 + } + fn uint64_value(i: usize) -> u64 { + i as u64 + (1 << 40) + } + fn float32_value(i: usize) -> f32 { + i as f32 * 0.5 + } + fn float64_value(i: usize) -> f64 { + i as f64 * -0.25 + } + fn bool_value(i: usize) -> bool { + i % 3 == 0 + } + fn string_value(i: usize) -> String { + format!("s{}", i % 101) + } + fn category_value(i: usize) -> &'static str { + ["red", "green", "blue"][i % 3] + } + #[cfg(feature = "datetime")] + fn date32_value(i: usize) -> i32 { + i as i32 * 3 + } + + fn column(name: &str, dtype: ArrowType, array: Array) -> FieldArray { + FieldArray::new(Field::new(name, dtype, true, None), array) + } + + fn all_types_table() -> Table { + let mask = null_mask(); + let strings: Vec = (0..N_ROWS).map(string_value).collect(); + let categories: Vec<&str> = (0..N_ROWS).map(category_value).collect(); + + let mut cols = vec![ + column( + "int32", + ArrowType::Int32, + Array::from_int32(IntegerArray::from_vec64( + (0..N_ROWS).map(int32_value).collect::>(), + Some(mask.clone()), + )), + ), + column( + "uint32", + ArrowType::UInt32, + Array::from_uint32(IntegerArray::from_vec64( + (0..N_ROWS).map(uint32_value).collect::>(), + Some(mask.clone()), + )), + ), + column( + "int64", + ArrowType::Int64, + Array::from_int64(IntegerArray::from_vec64( + (0..N_ROWS).map(int64_value).collect::>(), + Some(mask.clone()), + )), + ), + column( + "uint64", + ArrowType::UInt64, + Array::from_uint64(IntegerArray::from_vec64( + (0..N_ROWS).map(uint64_value).collect::>(), + Some(mask.clone()), + )), + ), + column( + "float32", + ArrowType::Float32, + Array::from_float32(FloatArray::from_vec64( + (0..N_ROWS).map(float32_value).collect::>(), + Some(mask.clone()), + )), + ), + column( + "float64", + ArrowType::Float64, + Array::from_float64(FloatArray::from_vec64( + (0..N_ROWS).map(float64_value).collect::>(), + Some(mask.clone()), + )), + ), + column( + "bool", + ArrowType::Boolean, + Array::from_bool(BooleanArray::new( + Bitmask::from_bools(&(0..N_ROWS).map(bool_value).collect::>()), + Some(mask.clone()), + )), + ), + column( + "string", + ArrowType::String, + Array::from_string32(StringArray::from_vec( + strings.iter().map(String::as_str).collect(), + Some(mask.clone()), + )), + ), + ]; + + #[cfg(feature = "default_categorical_8")] + cols.push(column( + "category", + ArrowType::Dictionary(CategoricalIndexType::UInt8), + Array::from_categorical8(CategoricalArray::::from_vec( + categories, + Some(mask.clone()), + )), + )); + #[cfg(not(feature = "default_categorical_8"))] + cols.push(column( + "category", + ArrowType::Dictionary(CategoricalIndexType::UInt32), + Array::from_categorical32(CategoricalArray::::from_vec( + categories, + Some(mask.clone()), + )), + )); + + #[cfg(feature = "datetime")] + { + // Date64 is left out: Parquet has no 64-bit DATE annotation, so + // the type mapping stores it as a plain INT64. + cols.push(column( + "date32", + ArrowType::Date32, + Array::from_datetime_i32(DatetimeArray::from_vec64( + (0..N_ROWS).map(date32_value).collect::>(), + Some(mask.clone()), + None, + )), + )); + } + + Table::new("all_types".to_string(), Some(cols)) + } + + fn roundtrip(table: &Table, compression: Option) -> Table { + let mut buf = Cursor::new(Vec::new()); + write_parquet_table(table, &mut buf, compression).expect("write"); + buf.seek(SeekFrom::Start(0)).unwrap(); + load_parquet_table(&mut buf).expect("read") + } + + fn col<'a>(table: &'a Table, name: &str) -> &'a FieldArray { + table + .cols + .iter() + .find(|c| c.field.name == name) + .unwrap_or_else(|| panic!("column {name} missing")) + } + + fn assert_all_types(out: &Table) { + let expected_nulls = (0..N_ROWS).filter(|&i| !is_valid(i)).count(); + assert_eq!(out.n_rows, N_ROWS); + for c in &out.cols { + assert!(c.field.nullable, "{} must be nullable", c.field.name); + assert_eq!(c.null_count, expected_nulls, "{} null count", c.field.name); + assert_eq!(c.array.len(), N_ROWS, "{} length", c.field.name); + } + + match &col(out, "int32").array { + Array::NumericArray(NumericArray::Int32(a)) => { + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(int32_value)) + } + other => panic!("int32: {other:?}"), + } + match &col(out, "uint32").array { + Array::NumericArray(NumericArray::UInt32(a)) => { + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(uint32_value)) + } + other => panic!("uint32: {other:?}"), + } + match &col(out, "int64").array { + Array::NumericArray(NumericArray::Int64(a)) => { + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(int64_value)) + } + other => panic!("int64: {other:?}"), + } + match &col(out, "uint64").array { + Array::NumericArray(NumericArray::UInt64(a)) => { + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(uint64_value)) + } + other => panic!("uint64: {other:?}"), + } + match &col(out, "float32").array { + Array::NumericArray(NumericArray::Float32(a)) => { + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(float32_value)) + } + other => panic!("float32: {other:?}"), + } + match &col(out, "float64").array { + Array::NumericArray(NumericArray::Float64(a)) => { + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(float64_value)) + } + other => panic!("float64: {other:?}"), + } + match &col(out, "bool").array { + Array::BooleanArray(a) => { + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(bool_value)) + } + other => panic!("bool: {other:?}"), + } + let strings: Vec> = match &col(out, "string").array { + Array::TextArray(TextArray::String32(a)) => { + (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect() + } + #[cfg(feature = "large_string")] + Array::TextArray(TextArray::String64(a)) => { + (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect() + } + other => panic!("string: {other:?}"), + }; + assert_eq!(strings, expected(string_value)); + + let categories: Vec> = match &col(out, "category").array { + #[cfg(feature = "default_categorical_8")] + Array::TextArray(TextArray::Categorical8(a)) => { + (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect() + } + #[cfg(not(feature = "default_categorical_8"))] + Array::TextArray(TextArray::Categorical32(a)) => { + (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect() + } + other => panic!("category: {other:?}"), + }; + assert_eq!(categories, expected(|i| category_value(i).to_owned())); + + #[cfg(feature = "datetime")] + { + assert_eq!(col(out, "date32").field.dtype, ArrowType::Date32); + match &col(out, "date32").array { + Array::TemporalArray(TemporalArray::Datetime32(a)) => assert_eq!( + (0..N_ROWS).map(|i| a.get(i)).collect::>(), + expected(date32_value) + ), + other => panic!("date32: {other:?}"), + } + } + } + + #[test] + fn nullable_all_types_uncompressed() { + let table = all_types_table(); + assert_all_types(&roundtrip(&table, None)); + } + + #[cfg(feature = "snappy")] + #[test] + fn nullable_all_types_snappy() { + let table = all_types_table(); + assert_all_types(&roundtrip(&table, Some(Compression::Snappy))); + } + + #[cfg(feature = "zstd")] + #[test] + fn nullable_all_types_zstd() { + let table = all_types_table(); + assert_all_types(&roundtrip(&table, Some(Compression::Zstd))); + } +} diff --git a/rust/tests/pyarrow_parquet.rs b/rust/tests/pyarrow_parquet.rs new file mode 100644 index 0000000..4517912 --- /dev/null +++ b/rust/tests/pyarrow_parquet.rs @@ -0,0 +1,219 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! Parquet reader conformance against files written by pyarrow. +//! +//! The fixtures under `pyarrow-roundtrip/` are committed and regenerated +//! with `generate_pyarrow_parquet_files.py`. They cover the page layouts +//! parquet-cpp writes by default (Snappy, dictionary encoding, DataPageV1, +//! multiple row groups) plus the DataPageV2 layouts with and without +//! dictionaries. + +#[cfg(feature = "parquet")] +mod pyarrow_parquet_tests { + use std::fs::File; + use std::io::BufReader; + + use lightstream::models::readers::parquet::{load_parquet_table, load_parquet_table_cols}; + use minarrow::{Array, ArrowType, MaskedArray, NumericArray, Table, TextArray}; + + fn open(name: &str) -> BufReader { + let path = format!("{}/pyarrow-roundtrip/{name}", env!("CARGO_MANIFEST_DIR")); + BufReader::new(File::open(&path).unwrap_or_else(|e| panic!("open {path}: {e}"))) + } + + fn column<'a>(table: &'a Table, name: &str) -> &'a Array { + let col = table + .cols + .iter() + .find(|c| c.field.name == name) + .unwrap_or_else(|| panic!("column {name} missing")); + &col.array + } + + fn int32_values(table: &Table, name: &str) -> Vec> { + match column(table, name) { + Array::NumericArray(NumericArray::Int32(a)) => (0..a.len()).map(|i| a.get(i)).collect(), + other => panic!("{name}: expected Int32, got {other:?}"), + } + } + + fn int64_values(table: &Table, name: &str) -> Vec> { + match column(table, name) { + Array::NumericArray(NumericArray::Int64(a)) => (0..a.len()).map(|i| a.get(i)).collect(), + other => panic!("{name}: expected Int64, got {other:?}"), + } + } + + fn float32_values(table: &Table, name: &str) -> Vec> { + match column(table, name) { + Array::NumericArray(NumericArray::Float32(a)) => { + (0..a.len()).map(|i| a.get(i)).collect() + } + other => panic!("{name}: expected Float32, got {other:?}"), + } + } + + fn float64_values(table: &Table, name: &str) -> Vec> { + match column(table, name) { + Array::NumericArray(NumericArray::Float64(a)) => { + (0..a.len()).map(|i| a.get(i)).collect() + } + other => panic!("{name}: expected Float64, got {other:?}"), + } + } + + fn bool_values(table: &Table, name: &str) -> Vec> { + match column(table, name) { + Array::BooleanArray(a) => (0..a.len()).map(|i| a.get(i)).collect(), + other => panic!("{name}: expected Boolean, got {other:?}"), + } + } + + /// UTF-8 columns map to `String`, or to `LargeString` when the + /// `large_string` feature is on. + fn string_values(table: &Table, name: &str) -> Vec> { + match column(table, name) { + Array::TextArray(TextArray::String32(a)) => { + (0..a.len()).map(|i| a.get(i).map(str::to_owned)).collect() + } + #[cfg(feature = "large_string")] + Array::TextArray(TextArray::String64(a)) => { + (0..a.len()).map(|i| a.get(i).map(str::to_owned)).collect() + } + other => panic!("{name}: expected a UTF-8 column, got {other:?}"), + } + } + + /// Assert the seven-row table written by `nullable_table()` in the + /// generator script, in schema order. + fn assert_nullable_table(table: &Table) { + assert_eq!(table.n_rows, 7); + assert_eq!(table.cols.len(), 6); + let names: Vec<&str> = table.cols.iter().map(|c| c.field.name.as_str()).collect(); + assert_eq!( + names, + ["int32", "int64", "float32", "float64", "bool", "string"] + ); + assert!(table.cols.iter().all(|c| c.field.nullable)); + + assert_eq!( + int32_values(table, "int32"), + [Some(1), None, Some(3), Some(4), None, Some(6), Some(7)] + ); + assert_eq!( + int64_values(table, "int64"), + [Some(100), Some(101), None, Some(103), Some(104), Some(105), None] + ); + assert_eq!( + float32_values(table, "float32"), + [Some(0.5), Some(1.5), Some(2.5), None, Some(4.5), Some(5.5), Some(6.5)] + ); + assert_eq!( + float64_values(table, "float64"), + [None, Some(-1.0), Some(-2.0), Some(-3.0), Some(-4.0), None, Some(-6.0)] + ); + assert_eq!( + bool_values(table, "bool"), + [Some(true), Some(false), None, Some(true), Some(false), Some(true), None] + ); + assert_eq!( + string_values(table, "string"), + [ + Some("a".into()), + Some("b".into()), + Some("a".into()), + None, + Some("c".into()), + Some("b".into()), + Some("a".into()) + ] + ); + for col in &table.cols { + assert_eq!(col.null_count, col.array.null_count(), "{}", col.field.name); + } + } + + /// The three-column file from the original defect report, written + /// with pyarrow's defaults: Snappy, dictionary encoding, DataPageV1 + /// and REQUIRED columns. + #[cfg(feature = "snappy")] + #[test] + fn reads_pyarrow_default_layout() { + let table = load_parquet_table(open("pyarrow_simple.parquet")).expect("read"); + assert_eq!(table.n_rows, 5); + assert_eq!(table.cols.len(), 3); + assert_eq!(table.cols[0].field.dtype, ArrowType::Int64); + #[cfg(not(feature = "large_string"))] + assert_eq!(table.cols[1].field.dtype, ArrowType::String); + #[cfg(feature = "large_string")] + assert_eq!(table.cols[1].field.dtype, ArrowType::LargeString); + assert_eq!(table.cols[2].field.dtype, ArrowType::Float64); + + assert_eq!( + int64_values(&table, "id"), + [Some(1), Some(2), Some(3), Some(4), Some(5)] + ); + assert_eq!( + string_values(&table, "name"), + ["Alice", "Bob", "Charlie", "Diana", "Eve"] + .map(|s| Some(s.to_owned())) + ); + assert_eq!( + float64_values(&table, "score"), + [Some(85.5), Some(92.0), Some(78.3), Some(95.1), Some(88.7)] + ); + assert!(table.cols.iter().all(|c| c.null_count == 0)); + } + + /// Nullable columns holding nulls, split over three row groups, with + /// pyarrow's default page layout. + #[cfg(feature = "snappy")] + #[test] + fn reads_pyarrow_nullable_columns_across_row_groups() { + let table = + load_parquet_table(open("pyarrow_nullable_row_groups.parquet")).expect("read"); + assert_nullable_table(&table); + } + + /// DataPageV2 pages with PLAIN values and no compression. + #[test] + fn reads_pyarrow_plain_data_page_v2() { + let table = load_parquet_table(open("pyarrow_plain_v2.parquet")).expect("read"); + assert_nullable_table(&table); + } + + /// DataPageV2 pages with dictionary-encoded values and Snappy + /// compression on the value section only. + #[cfg(feature = "snappy")] + #[test] + fn reads_pyarrow_dictionary_data_page_v2() { + let table = load_parquet_table(open("pyarrow_dictionary_v2.parquet")).expect("read"); + assert_nullable_table(&table); + } + + /// Column projection over a pyarrow file selects by name and keeps + /// the values of the selected columns intact. + #[cfg(feature = "snappy")] + #[test] + fn projects_columns_from_pyarrow_file() { + let table = + load_parquet_table_cols(open("pyarrow_simple.parquet"), &["score", "name"]) + .expect("read"); + assert_eq!(table.n_rows, 5); + let names: Vec<&str> = table.cols.iter().map(|c| c.field.name.as_str()).collect(); + assert_eq!(names, ["name", "score"]); + assert_eq!( + string_values(&table, "name"), + ["Alice", "Bob", "Charlie", "Diana", "Eve"] + .map(|s| Some(s.to_owned())) + ); + assert_eq!( + float64_values(&table, "score"), + [Some(85.5), Some(92.0), Some(78.3), Some(95.1), Some(88.7)] + ); + } +} From 6b7be7330d01a0347ffc1a2eaf3232e8a7ec086f Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Thu, 10 Sep 2026 05:47:22 +0100 Subject: [PATCH 4/7] Update CHANGELOG.md --- CHANGELOG.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 235d1ae..a77c8af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,20 +6,20 @@ Notable changes are recorded from 0.5.0 onward. ### Changed -- minarrow 0.18.1 and vec64 0.5.1. +- Minarrow upgraded to 0.18.1 and vec64 0.5.1. - New `decimal` feature forwarding minarrow's `decimal` feature. Decimal32, Decimal64 and Decimal128 columns are supported in Arrow IPC and in Parquet, where they map to the DECIMAL logical type over INT32, INT64 and FIXED_LEN_BYTE_ARRAY. - The Python package pins minarrow and minarrow-pyo3 at 0.18.1. - The Parquet writer follows the Parquet value layout for nullable columns. Value sections hold non-null values only and DataPageV2 headers count every row in `num_values`, so files with nulls now read in pyarrow and other Parquet readers. Files with nulls written by earlier releases do not read back under this release. ### Fixed -- The Parquet reader failed with `UnexpectedEof` on files written by pyarrow. It now reads DataPageV1 pages with compressed levels, `RLE` booleans, dictionary-encoded columns of any physical type, and files with several row groups. Dictionary pages in files without lightstream's `created_by` marker expand to the schema type rather than being read as categorical columns. +- The Parquet reader failed with `UnexpectedEof` on files written by pyarrow. It now reads DataPageV1 pages with compressed levels, `RLE` booleans, dictionary-encoded columns of any physical type, and files with several row groups. Dictionary pages in files without lightstream's `created_by` marker expand to the schema type instead of being read as categorical columns. ## 0.6.1 ### Changed -- minarrow 0.17.0, vec64 0.5.0, arrow 59.2.0 and polars 0.55.2. +- minarrow version bump to 0.17.0, vec64 0.5.0, arrow 59.2.0 and polars 0.55.2. ## 0.6.0 From 294bfd7796256ed202c79d67c83ec0ba3bad69ae Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Thu, 10 Sep 2026 06:32:48 +0100 Subject: [PATCH 5/7] Map Parquet columns by schema type with LogicalType support, fixed-width decimals and unit-correct temporal storage --- CHANGELOG.md | 5 +- python/tests/test_files.py | 6 + rust/Cargo.lock | 2 + .../generate_pyarrow_parquet_files.py | 35 ++ .../pyarrow_temporal_decimal.parquet | Bin 0 -> 3054 bytes rust/src/constants.rs | 3 - rust/src/models/encoders/parquet/data.rs | 45 ++ rust/src/models/encoders/parquet/metadata.rs | 86 ++- rust/src/models/readers/chunked/parquet.rs | 19 +- rust/src/models/readers/parquet.rs | 524 +++++++++--------- rust/src/models/types/parquet.rs | 111 ++-- rust/src/models/writers/parquet.rs | 50 +- rust/tests/parquet_nullable_roundtrip.rs | 307 ++++++++-- rust/tests/pyarrow_parquet.rs | 125 ++++- 14 files changed, 932 insertions(+), 386 deletions(-) create mode 100644 rust/pyarrow-roundtrip/pyarrow_temporal_decimal.parquet diff --git a/CHANGELOG.md b/CHANGELOG.md index 235d1ae..ff1f16e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,10 +10,13 @@ Notable changes are recorded from 0.5.0 onward. - New `decimal` feature forwarding minarrow's `decimal` feature. Decimal32, Decimal64 and Decimal128 columns are supported in Arrow IPC and in Parquet, where they map to the DECIMAL logical type over INT32, INT64 and FIXED_LEN_BYTE_ARRAY. - The Python package pins minarrow and minarrow-pyo3 at 0.18.1. - The Parquet writer follows the Parquet value layout for nullable columns. Value sections hold non-null values only and DataPageV2 headers count every row in `num_values`, so files with nulls now read in pyarrow and other Parquet readers. Files with nulls written by earlier releases do not read back under this release. +- Parquet column types follow the schema element. Categorical columns are UTF8 string columns with dictionary-encoded pages and read back as strings. Date64 is stored as DATE days, seconds-unit timestamps and times as milliseconds, and TIME columns take the INT32 or INT64 width their unit requires. UTF8 columns read as `String` and widen to `LargeString` only when the data exceeds 32-bit offsets. +- The `LogicalType` schema annotation is written and read, so nanosecond timestamps and times round-trip and pyarrow nanosecond timestamps no longer read as Int64. +- Duration and Interval columns are reported as `UnsupportedType` by the Parquet writer instead of panicking, and INT96 columns are reported by name on read. ### Fixed -- The Parquet reader failed with `UnexpectedEof` on files written by pyarrow. It now reads DataPageV1 pages with compressed levels, `RLE` booleans, dictionary-encoded columns of any physical type, and files with several row groups. Dictionary pages in files without lightstream's `created_by` marker expand to the schema type rather than being read as categorical columns. +- The Parquet reader failed with `UnexpectedEof` on files written by pyarrow. It now reads DataPageV1 pages with compressed levels, `RLE` booleans, dictionary-encoded columns of any physical type, and files with several row groups. ## 0.6.1 diff --git a/python/tests/test_files.py b/python/tests/test_files.py index af8d9ae..575bfd0 100644 --- a/python/tests/test_files.py +++ b/python/tests/test_files.py @@ -11,6 +11,7 @@ """ import gc +from decimal import Decimal import lightstream as ls import minarrow @@ -236,6 +237,11 @@ def nullable_table(): "ratio": pa.array([0.5, 1.5, 2.5, None, 4.5], type=pa.float32()), "flag": pa.array([True, False, None, True, False], type=pa.bool_()), "day": pa.array([1, None, 3, 4, 5], type=pa.date32()), + "at": pa.array([1_000, 2_000, None, 4_000, 5_000], type=pa.timestamp("ns")), + "amount": pa.array( + [Decimal("1.25"), None, Decimal("-3.50"), Decimal("0.01"), Decimal("99.99")], + type=pa.decimal128(10, 2), + ), } ) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 9238e94..caccf56 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2049,6 +2049,8 @@ checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" [[package]] name = "minarrow" version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e53f5031b7d16aa76bc8671171814419694c687f5160917a3668f010dab2b609" dependencies = [ "arrow", "arrow-schema", diff --git a/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py b/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py index be5530b..00d5178 100644 --- a/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py +++ b/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py @@ -13,6 +13,9 @@ encoding off, DataPageV2 and no compression - pyarrow_dictionary_v2.parquet: the same nullable data with dictionary encoding on, DataPageV2 and Snappy + - pyarrow_temporal_decimal.parquet: date, time, timestamp and decimal + columns with nulls, written with pyarrow's defaults. Nanosecond units + and decimals exercise the LogicalType annotation The fixtures are committed alongside this script. Run from rust/ to regenerate them: @@ -21,6 +24,7 @@ """ import os +from decimal import Decimal import pyarrow as pa import pyarrow.parquet as pq @@ -53,6 +57,33 @@ def nullable_table() -> pa.Table: ) +def temporal_decimal_table() -> pa.Table: + """Five rows of temporal and decimal types, matching the expectations in + reads_pyarrow_temporal_and_decimal_columns in Rust.""" + return pa.table( + { + "date": pa.array([0, 1, None, 19_000, -5], type=pa.date32()), + "time_ms": pa.array([0, 1_000, None, 43_200_000, 86_399_999], type=pa.time32("ms")), + "time_us": pa.array([0, None, 2_000_000, 43_200_000_000, 86_399_999_999], type=pa.time64("us")), + "ts_ms": pa.array([0, 1_700_000_000_000, None, -1, 86_400_000], type=pa.timestamp("ms")), + "ts_us": pa.array([0, 1_700_000_000_000_000, None, -1, 1], type=pa.timestamp("us")), + "ts_ns": pa.array([0, 1_700_000_000_000_000_000, None, -1, 1], type=pa.timestamp("ns")), + "dec32": pa.array( + [Decimal("1.25"), None, Decimal("-3.50"), Decimal("0.01"), Decimal("99999.99")], + type=pa.decimal32(7, 2), + ), + "dec64": pa.array( + [Decimal("1.2345"), Decimal("-1.0000"), None, Decimal("0.0001"), Decimal("12345678901234.5678")], + type=pa.decimal64(18, 4), + ), + "dec128": pa.array( + [Decimal("1.000001"), None, Decimal("-1.000001"), Decimal("0.000000"), Decimal("12345678901234567890123456789012.123456")], + type=pa.decimal128(38, 6), + ), + } + ) + + def main() -> None: simple = simple_table() pq.write_table(simple, os.path.join(OUT_DIR, "pyarrow_simple.parquet")) @@ -76,6 +107,10 @@ def main() -> None: data_page_version="2.0", compression="snappy", ) + pq.write_table( + temporal_decimal_table(), + os.path.join(OUT_DIR, "pyarrow_temporal_decimal.parquet"), + ) for name in sorted(os.listdir(OUT_DIR)): if name.startswith("pyarrow_") and name.endswith(".parquet"): print(f"wrote {name}") diff --git a/rust/pyarrow-roundtrip/pyarrow_temporal_decimal.parquet b/rust/pyarrow-roundtrip/pyarrow_temporal_decimal.parquet new file mode 100644 index 0000000000000000000000000000000000000000..5a8bef6c49ff2827c39d617b4bfe333d3946b2bb GIT binary patch literal 3054 zcmb7GT})eL82-+oEr-fVLB4t#yV(mknaU_-fZc)#U!g+Bk5F;^xMX3amW}cwg`q^_ zF1#?&L=zJ=kt~ZB#*4))#_Y~{vl#EZF>x1WS1ws}F(eu1NuZ=Jy@2Qn_eUV6`+4vAWR8tO-Y~T$z5VZ-4ir z$87hjkTO&A5UC5TM0)rszELXee~Qo)_VI^ry5jQ?s(VO>-b?x9G+nNpGaj#^b)71C z70;bLZ_{xQC^tVD{jBTv2^WcGOYC^=>{035sFTz#-4%sT+F2aH!&g{iQ&N<$YI|RQ z)A_Xb=Z1P^QILAwfag(Y}YgO<}nD=p#cy-fENvrD;^ z1>VBA+OF!J#f8mWd*fJv-V_1#n-^Dd=U0k_*IObHfPg&*@Ro!-%CL8{2;KIUv3LKa zP4HH+y_MiY+N;F^K-qd251`J=xvvM-K^foTHgo1872~tgi<$jkiID*_1^g z^Wk2+I{{%MAA=Axgdu-865cgfgsb6R<_m_+7`YnqhrUokHx*;V(5kn#a0i?~I_O&l zCs_ZJzObQE4{@5cGv%K9` zJ;>fKit%=X@fE6$8rIyfs%I}?Rmyz4jyG`EZ(Q}`4U{_+y@B%Sp)oiu_x0G>LJXfb z8b3UO7~JiMvgJl=X0pCWl4?&4-Byk30pr$@h+adV0d);pRb2^P3x@oL7W%7+g(G^S z$Yh?MIc>Hn{ed?|-_N9OXJgNx(hz| zC0D7sLUQ;Ck1stPizTp0jey+($q{?JNZ|7P%y{<%LF3OL=)?=Om3&fjDEkn9kPpkB zoiybIK}=CTke4+QjnE{cpll+3N-fsp;F5W!*%6<#w>UEK@`Y8I+EGh?NTjC8KWL{$ ztHjfljh~8Vx1lz4!O~BQ>?F2?O|8jP%OqM^2P%?;j%0RbyW=}_o{^(K7fe5LEWh=LNBD=nRVqs%rZ9957 fw^6*fwi@m3?d) { @@ -168,6 +169,50 @@ pub fn encode_large_string_plain( Ok(()) } +// Temporal values + +/// Plain-encode temporal values in the unit of the column's Parquet +/// annotation. +/// +/// `mul` and `div` come from `temporal_unit_scale` and carry each value +/// from its Arrow unit into the Parquet unit. `physical` selects the +/// storage width: INT32 for DATE and TIME(MILLIS), INT64 for the other +/// units. A value outside the INT32 range is an `InputDataError`. +pub(crate) fn encode_temporal_plain>( + data: &[T], + mul: i64, + div: i64, + physical: ParquetPhysicalType, + out: &mut Vec, +) -> Result<(), IoError> { + match physical { + ParquetPhysicalType::Int32 => { + out.reserve(data.len() * 4); + for &v in data { + let scaled = v.into() * mul / div; + let narrowed = i32::try_from(scaled).map_err(|_| { + IoError::InputDataError(format!( + "temporal value {scaled} does not fit the INT32 storage of its Parquet type" + )) + })?; + out.extend_from_slice(&narrowed.to_le_bytes()); + } + } + ParquetPhysicalType::Int64 => { + out.reserve(data.len() * 8); + for &v in data { + out.extend_from_slice(&(v.into() * mul / div).to_le_bytes()); + } + } + other => { + return Err(IoError::Internal(format!( + "temporal column mapped to physical type {other:?}" + ))); + } + } + Ok(()) +} + // Temporal aliases for the same physical type /// Encode `DATE32`/`TIME32` using Parquet plain format. diff --git a/rust/src/models/encoders/parquet/metadata.rs b/rust/src/models/encoders/parquet/metadata.rs index 61075fc..045ee73 100644 --- a/rust/src/models/encoders/parquet/metadata.rs +++ b/rust/src/models/encoders/parquet/metadata.rs @@ -18,7 +18,7 @@ use std::io::{Seek, Write}; use crate::constants::PARQUET_MAGIC; use crate::error::IoError; -use crate::models::types::parquet::{ParquetEncoding, ParquetPhysicalType}; +use crate::models::types::parquet::{ParquetEncoding, ParquetLogicalType, ParquetPhysicalType}; // --------------------- Structs ------------------------------------ // @@ -50,6 +50,10 @@ pub(crate) struct SchemaElement { pub type_: Option, /// Legacy converted type ID (if any). pub converted_type: Option, + /// `LogicalType` annotation (field 10). Written alongside the converted + /// type so readers of either generation see the annotation. Nanosecond + /// units exist only here. + pub logical_type: Option, /// Type length (e.g. for FIXED_LEN_BYTE_ARRAY). pub type_length: Option, /// Decimal precision (if applicable). @@ -436,11 +440,85 @@ impl SchemaElement { thrift_write_field_i32(&mut w, &mut last, 9, id); } + // [10] LogicalType - a union struct whose single set field selects + // the annotation. Empty structs mark STRING and DATE. TIME and + // TIMESTAMP carry isAdjustedToUTC and a TimeUnit union. INTEGER + // carries the bit width and sign. DECIMAL carries scale then + // precision. + if let Some(logical) = &self.logical_type + && let Some(union_id) = logical_type_union_id(logical) + { + thrift_write_field_struct_begin(&mut w, &mut last, 10); + let mut union_last = 0i16; + thrift_write_field_struct_begin(&mut w, &mut union_last, union_id); + let mut inner = 0i16; + match logical { + #[cfg(feature = "datetime")] + ParquetLogicalType::TimestampMillis + | ParquetLogicalType::TimestampMicros + | ParquetLogicalType::TimestampNanos + | ParquetLogicalType::TimeMillis + | ParquetLogicalType::TimeMicros + | ParquetLogicalType::TimeNanos => { + thrift_write_field_bool(&mut w, &mut inner, 1, false); + thrift_write_field_struct_begin(&mut w, &mut inner, 2); + let mut unit_last = 0i16; + thrift_write_field_struct_begin(&mut w, &mut unit_last, time_unit_union_id(logical)); + thrift_write_field_stop(&mut w); + thrift_write_field_stop(&mut w); + } + ParquetLogicalType::IntType { bit_width, is_signed } => { + thrift_write_field_i8(&mut w, &mut inner, 1, *bit_width as i8); + thrift_write_field_bool(&mut w, &mut inner, 2, *is_signed); + } + #[cfg(feature = "decimal")] + ParquetLogicalType::Decimal { precision, scale } => { + thrift_write_field_i32(&mut w, &mut inner, 1, *scale as i32); + thrift_write_field_i32(&mut w, &mut inner, 2, *precision as i32); + } + _ => {} + } + thrift_write_field_stop(&mut w); + thrift_write_field_stop(&mut w); + } + thrift_write_field_stop(&mut w); Ok(()) } } +/// Field id of the `LogicalType` union member for a logical annotation. +/// `None` for annotations with no `LogicalType` form. +fn logical_type_union_id(logical: &ParquetLogicalType) -> Option { + Some(match logical { + ParquetLogicalType::NoneType => return None, + ParquetLogicalType::Utf8 => 1, + #[cfg(feature = "decimal")] + ParquetLogicalType::Decimal { .. } => 5, + #[cfg(feature = "datetime")] + ParquetLogicalType::Date32 => 6, + #[cfg(feature = "datetime")] + ParquetLogicalType::TimeMillis + | ParquetLogicalType::TimeMicros + | ParquetLogicalType::TimeNanos => 7, + #[cfg(feature = "datetime")] + ParquetLogicalType::TimestampMillis + | ParquetLogicalType::TimestampMicros + | ParquetLogicalType::TimestampNanos => 8, + ParquetLogicalType::IntType { .. } => 10, + }) +} + +/// Field id of the `TimeUnit` union member: MILLIS 1, MICROS 2, NANOS 3. +#[cfg(feature = "datetime")] +fn time_unit_union_id(logical: &ParquetLogicalType) -> i16 { + match logical { + ParquetLogicalType::TimestampMillis | ParquetLogicalType::TimeMillis => 1, + ParquetLogicalType::TimestampMicros | ParquetLogicalType::TimeMicros => 2, + _ => 3, + } +} + impl RowGroupMeta { /// Write a row group descriptor via TCompactProtocol. pub fn write(&self, mut w: W) -> Result<(), IoError> { @@ -573,6 +651,7 @@ impl Statistics { /// Compact-protocol element and field type identifiers from the Thrift spec. const TC_BOOL_TRUE: u8 = 1; const TC_BOOL_FALSE: u8 = 2; +const TC_BYTE: u8 = 3; const TC_I32: u8 = 5; const TC_I64: u8 = 6; const TC_BINARY: u8 = 8; @@ -635,6 +714,11 @@ fn thrift_write_list_header(w: &mut W, elem_type: u8, len: usize) { fn thrift_write_field_stop(w: &mut W) { w.write_all(&[0]).unwrap(); } +/// Write a byte field as one raw byte. +fn thrift_write_field_i8(w: &mut W, last: &mut i16, id: i16, v: i8) { + thrift_write_field_header(w, last, id, TC_BYTE); + w.write_all(&[v as u8]).unwrap(); +} /// Write an i32 field as a zigzag varint. fn thrift_write_field_i32(w: &mut W, last: &mut i16, id: i16, v: i32) { thrift_write_field_header(w, last, id, TC_I32); diff --git a/rust/src/models/readers/chunked/parquet.rs b/rust/src/models/readers/chunked/parquet.rs index 6b80552..61a4505 100644 --- a/rust/src/models/readers/chunked/parquet.rs +++ b/rust/src/models/readers/chunked/parquet.rs @@ -156,7 +156,7 @@ mod tests { } #[test] - fn categorical_column_roundtrips_via_dictionary_page() { + fn categorical_column_reads_back_as_utf8_strings() { // Regression: the reader's DictionaryPageHeader inner parser used // to consume raw i32 values without the Thrift type/id prefixes // the writer emits. That left the cursor 3 bytes shy of the @@ -164,11 +164,15 @@ mod tests { // length prefix was parsed as garbage - causing // `parse_dictionary_values` to ask for ~50 MB and fail with // `UnexpectedEof`. This test exercises the round-trip end-to-end. + // + // Parquet has no dictionary type, so the categorical column is a + // UTF8 string column with dictionary-encoded pages and reads back + // as strings. use crate::models::readers::parquet::load_parquet_table; use crate::models::writers::parquet::write_parquet_table; use minarrow::{ - Array, ArrowType, Bitmask, Buffer, CategoricalArray, Field, FieldArray, TextArray, - Vec64, ffi::arrow_dtype::CategoricalIndexType, + Array, ArrowType, Bitmask, Buffer, CategoricalArray, Field, FieldArray, MaskedArray, + TextArray, Vec64, ffi::arrow_dtype::CategoricalIndexType, }; use std::sync::Arc; @@ -228,6 +232,15 @@ mod tests { .expect("categorical column must round-trip via Parquet"); assert_eq!(got.n_rows, n_rows); assert_eq!(got.cols.len(), 1); + assert_eq!(got.cols[0].field.dtype, ArrowType::String); + match &got.cols[0].array { + Array::TextArray(TextArray::String32(strings)) => { + let expected: Vec> = + (0..n_rows).map(|i| Some(["red", "green", "blue"][i % 3])).collect(); + assert_eq!((0..n_rows).map(|i| strings.get(i)).collect::>(), expected); + } + other => panic!("expected a String column, got {other:?}"), + } let _ = std::fs::remove_file(&path); } diff --git a/rust/src/models/readers/parquet.rs b/rust/src/models/readers/parquet.rs index 893b4c5..a33f6f1 100644 --- a/rust/src/models/readers/parquet.rs +++ b/rust/src/models/readers/parquet.rs @@ -22,14 +22,14 @@ //! - Works with any `Read + Seek` //! - Reads into memory - no mmap zero-copy like IPC at the present time. //! -//! ## Categorical columns -//! Parquet has no dictionary logical type. lightstream writes a categorical -//! column as a BYTE_ARRAY UTF8 leaf whose pages are all dictionary-encoded -//! and marks the file with [`PARQUET_CREATED_BY`](crate::constants::PARQUET_CREATED_BY). -//! Only files carrying that marker read such columns back as -//! `ArrowType::Dictionary`. Dictionary pages in other files are a storage -//! encoding and expand to the schema type, so a pyarrow string column reads -//! as a string column whether or not pyarrow dictionary-encoded it. +//! ## Types +//! The schema element decides the column type: its physical type plus the +//! `LogicalType` annotation, or the legacy converted type when that is all +//! the writer recorded. Page encodings do not take part, so a +//! dictionary-encoded column reads as the type its schema element names. +//! Categorical columns written by lightstream therefore come back as UTF8 +//! string columns. UTF8 columns use 32-bit offsets and widen to +//! `LargeString` only when the column's data outgrows them. //! //! ## Outputs //! On success returns a fully materialised `Table`; otherwise yields an `IOError` @@ -41,7 +41,7 @@ use std::io::{Cursor, Read, Seek, SeekFrom}; use std::sync::Arc; use crate::compression::{Compression, decompress}; -use crate::constants::{PARQUET_CREATED_BY, PARQUET_MAGIC}; +use crate::constants::PARQUET_MAGIC; use crate::error::IoError; #[cfg(feature = "datetime")] use crate::models::decoders::parquet::{decode_datetime32_plain, decode_datetime64_plain}; @@ -57,19 +57,25 @@ use crate::models::encoders::parquet::metadata::{ use crate::models::types::parquet::{ ParquetEncoding, ParquetLogicalType, ParquetPhysicalType, parquet_to_arrow_type, }; -use minarrow::ffi::arrow_dtype::CategoricalIndexType; use minarrow::{ - Array, ArrowType, Bitmask, BooleanArray, CategoricalArray, Field, FieldArray, FloatArray, - IntegerArray, NumericArray, StringArray, Table, TextArray, Vec64, vec64, + Array, ArrowType, Bitmask, BooleanArray, Field, FieldArray, FloatArray, IntegerArray, + NumericArray, StringArray, Table, TextArray, Vec64, vec64, }; #[cfg(feature = "decimal")] use minarrow::DecimalArray; #[cfg(feature = "datetime")] -use minarrow::{DatetimeArray, TemporalArray}; +use minarrow::{DatetimeArray, TemporalArray, TimeUnit}; -/// Build the logical type for a schema element, incorporating precision -/// and scale for DECIMAL columns where `from_converted_type` returns `None`. +/// Build the logical type for a schema element. +/// +/// The `LogicalType` annotation wins when the writer recorded one, since it +/// carries units the converted type cannot, such as nanoseconds. Otherwise +/// the converted type applies, with DECIMAL taking its precision and scale +/// from the schema element's own fields. fn logical_type_from_schema(se: &SchemaElement) -> Option { + if let Some(logical) = &se.logical_type { + return Some(logical.clone()); + } #[cfg(feature = "decimal")] if se.converted_type == Some(5) { let precision = se.precision.unwrap_or(0) as u8; @@ -79,19 +85,37 @@ fn logical_type_from_schema(se: &SchemaElement) -> Option { ParquetLogicalType::from_converted_type(se.converted_type) } -/// Decode a FIXED_LEN_BYTE_ARRAY Decimal128 buffer (16 big-endian bytes per -/// value) into a `Vec64`. +/// Decode a FIXED_LEN_BYTE_ARRAY decimal buffer into unscaled values. +/// +/// Each value is `width` bytes of big-endian two's complement, up to 16 +/// bytes. Values are sign-extended into `i128` and then narrowed to the +/// decimal width the schema's precision selects. #[cfg(feature = "decimal")] -fn decode_decimal128_plain(buf: &[u8]) -> Result, IoError> { - if buf.len() % 16 != 0 { - return Err(IoError::Format( - "decode_decimal128_plain: buffer len % 16 != 0".into(), - )); +fn decode_decimal_fixed_plain>( + buf: &[u8], + width: usize, +) -> Result, IoError> { + if width == 0 || width > 16 { + return Err(IoError::Format(format!( + "decimal FIXED_LEN_BYTE_ARRAY width {width} is outside 1..=16" + ))); + } + if buf.len() % width != 0 { + return Err(IoError::Format(format!( + "decimal value section of {} bytes is not a multiple of width {width}", + buf.len() + ))); } - Ok(buf - .chunks_exact(16) - .map(|c| i128::from_be_bytes(c.try_into().unwrap())) - .collect()) + buf.chunks_exact(width) + .map(|chunk| { + let fill = if chunk[0] & 0x80 != 0 { 0xff } else { 0x00 }; + let mut bytes = [fill; 16]; + bytes[16 - width..].copy_from_slice(chunk); + T::try_from(i128::from_be_bytes(bytes)).map_err(|_| { + IoError::Format("decimal value exceeds the width its precision allows".into()) + }) + }) + .collect() } /// Read an entire in-memory Table from a Parquet v2 file. @@ -106,27 +130,6 @@ pub fn load_parquet_table(r: R) -> Result { read_parquet_impl(r, None) } -/// Index type assumed for dictionary-encoded columns whose original -/// Arrow type was lost in the Parquet schema. -/// -/// The writer maps every `ArrowType::Dictionary(_)` to physical Int32 -/// with `NoneType` logical type, so the schema doesn't carry the index -/// width. Reads pick whichever width is the build's default categorical -/// type. Round-tripping a column written with `default_categorical_8` -/// disabled into a build that has it enabled (or vice versa) is not -/// supported. -#[inline] -fn default_categorical_index_type() -> CategoricalIndexType { - #[cfg(feature = "default_categorical_8")] - { - CategoricalIndexType::UInt8 - } - #[cfg(not(feature = "default_categorical_8"))] - { - CategoricalIndexType::UInt32 - } -} - /// Read only the named columns from a Parquet v2 file. /// /// Column names must match the schema's `path_in_schema` entries. Returns @@ -158,12 +161,7 @@ pub fn load_parquet_table_cols( /// receives one entry per row. /// - Dictionary-encoded pages of any physical type are expanded through /// the row group's dictionary, so dictionary encoding stays a storage -/// detail of the writer. -/// - A BYTE_ARRAY UTF8 leaf with a dictionary page in a file carrying -/// [`PARQUET_CREATED_BY`](crate::constants::PARQUET_CREATED_BY) is the -/// lightstream categorical convention and reads back as -/// `ArrowType::Dictionary`. Dictionaries from several row groups merge -/// into one set of unique values. +/// detail of the writer and the column keeps its schema type. /// - Repeated (nested) columns and unknown compression codecs are /// reported as errors rather than decoded. fn read_parquet_impl( @@ -222,10 +220,6 @@ fn read_parquet_impl( ))); } - // The categorical convention only applies to lightstream's own files. - // Other writers dictionary-encode any column type as a storage detail. - let lightstream_written = meta.created_by.as_deref() == Some(PARQUET_CREATED_BY); - let mut columns = Vec::with_capacity(leaves.len()); for (col_idx, &(leaf, physical)) in leaves.iter().enumerate() { @@ -249,42 +243,19 @@ fn read_parquet_impl( } }; - let logical = logical_type_from_schema(leaf); - let schema_ty = parquet_to_arrow_type(physical, logical)?; - let has_dictionary = meta - .row_groups - .iter() - .any(|rg| rg.columns[col_idx].meta_data.dictionary_page_offset.is_some()); - let categorical = - lightstream_written && physical == ParquetPhysicalType::ByteArray && has_dictionary; - let ty = if categorical { - ArrowType::Dictionary(default_categorical_index_type()) - } else { - schema_ty - }; - // Categorical columns accumulate their dictionary indices as PLAIN - // u32 entries. Every other column accumulates its physical values. - let layout = if categorical { - ValueLayout::Fixed(4) - } else { - ValueLayout::of(physical, leaf.type_length)? - }; + let ty = parquet_to_arrow_type(physical, logical_type_from_schema(leaf))?; + let layout = ValueLayout::of(physical, leaf.type_length)?; let mut def_levels: Vec = Vec::new(); let mut values: Vec = Vec::new(); - // Merged categorical dictionary across row groups, with the lookup - // used to remap each row group's local indices onto it. - let mut unique_values: Vec> = Vec::new(); - let mut unique_index: BTreeMap, u32> = BTreeMap::new(); for rg in &meta.row_groups { let cmeta = &rg.columns[col_idx].meta_data; let codec = map_codec(cmeta.codec)?; // The row group's dictionary entries without their PLAIN length - // prefix, plus the remap onto the merged categorical dictionary. + // prefix. let mut dict_entries: Vec> = Vec::new(); - let mut dict_remap: Vec = Vec::new(); if let Some(dict_off) = cmeta.dictionary_page_offset { r.seek(SeekFrom::Start(dict_off as u64))?; let ph = parse_page_header(&mut r)?; @@ -297,20 +268,7 @@ fn read_parquet_impl( Some(c) => decompress(&body, c)?, None => body, }; - if categorical { - for entry in parse_dictionary_values(&body)? { - let merged = match unique_index.get(&entry) { - Some(&idx) => idx, - None => { - let idx = unique_values.len() as u32; - unique_index.insert(entry.clone(), idx); - unique_values.push(entry); - idx - } - }; - dict_remap.push(merged); - } - } else { + { dict_entries = match layout { ValueLayout::LengthPrefixed => parse_dictionary_values(&body)?, ValueLayout::Fixed(width) => { @@ -374,13 +332,6 @@ fn read_parquet_impl( }; let mut plain = Vec::new(); for &idx in indices.iter() { - if categorical { - let merged = dict_remap.get(idx as usize).ok_or_else(|| { - IoError::Format(format!("dictionary index {idx} out of range")) - })?; - plain.extend_from_slice(&merged.to_le_bytes()); - continue; - } let entry = dict_entries.get(idx as usize).ok_or_else(|| { IoError::Format(format!("dictionary index {idx} out of range")) })?; @@ -414,8 +365,23 @@ fn read_parquet_impl( ))); } + // UTF8 columns use 32-bit offsets unless the data outgrows them. + #[cfg(feature = "large_string")] + let ty = if ty == ArrowType::String && values.len() > u32::MAX as usize { + ArrowType::LargeString + } else { + ty + }; + #[cfg(not(feature = "large_string"))] + if ty == ArrowType::String && values.len() > u32::MAX as usize { + return Err(IoError::UnsupportedType(format!( + "column '{}' holds more than 4 GiB of string data, which needs the large_string feature", + leaf.name + ))); + } + let null_count = def_levels.iter().filter(|&&b| !b).count(); - let array = decode_column(&ty, &unique_values, &values, def_levels.len(), def_levels)?; + let array = decode_column(&ty, physical, layout, &values, def_levels.len(), def_levels)?; columns.push(FieldArray { field: Field { @@ -438,6 +404,13 @@ fn read_parquet_impl( }) } +/// Error for a DECIMAL column stored in a physical type the reader does +/// not decode, which is BYTE_ARRAY. +#[cfg(feature = "decimal")] +fn decimal_storage_error(physical: ParquetPhysicalType) -> IoError { + IoError::UnsupportedType(format!("DECIMAL stored as {physical:?}")) +} + /// Byte layout of one value inside a PLAIN value section. #[derive(Clone, Copy, PartialEq, Eq, Debug)] enum ValueLayout { @@ -633,11 +606,14 @@ fn read_data_page_v2( /// Build the column array from PLAIN entries, one per row. /// /// `buf` holds one entry per row in the layout the PLAIN decoders expect, -/// with a zero entry under each null. Categorical columns pass their -/// merged dictionary in `dict` and hold u32 indices in `buf`. +/// with a zero entry under each null. `physical` and `layout` describe the +/// entries for types with more than one storage form, which for now means +/// decimals stored as INT32, INT64 or FIXED_LEN_BYTE_ARRAY. +#[cfg_attr(not(feature = "decimal"), allow(unused_variables))] fn decode_column( ty: &ArrowType, - dict: &[Vec], + physical: ParquetPhysicalType, + layout: ValueLayout, buf: &[u8], len: usize, def_levels: Vec, @@ -692,58 +668,42 @@ fn decode_column( }))) } - // dictionary / categoricals, u32 indices per row - ArrowType::Dictionary(key_ty) => match key_ty { - #[cfg(any( - not(feature = "default_categorical_8"), - feature = "extended_categorical" - ))] - CategoricalIndexType::UInt32 => { - build_cat32(decode_uint32_as_int32_plain(buf)?, dict, mask) - } - #[cfg(feature = "default_categorical_8")] - CategoricalIndexType::UInt8 => build_cat8(decode_uint32_as_int32_plain(buf)?, dict, mask), - #[cfg(all(feature = "extended_categorical", feature = "large_string"))] - CategoricalIndexType::UInt64 => { - let idx = decode_uint32_as_int32_plain(buf)? - .into_iter() - .map(|v| v as u64) - .collect(); - build_cat64(idx, dict, mask) - } - // Which index widths exist depends on minarrow's categorical - // feature flags, so this arm is unreachable in some builds. - #[allow(unreachable_patterns)] - _ => { - return Err(IoError::UnsupportedType(format!( - "dictionary index {:?}", - key_ty - ))); - } - }, - - // temporal + // temporal. DATE and TIME(MILLIS) are INT32, every other unit INT64. #[cfg(feature = "datetime")] ArrowType::Date32 => Array::TemporalArray(TemporalArray::Datetime32(Arc::new( DatetimeArray { data: decode_datetime32_plain(buf)?.into(), null_mask: mask, - time_unit: Default::default(), + time_unit: TimeUnit::Days, }, ))), #[cfg(feature = "datetime")] - ArrowType::Date64 => Array::TemporalArray(TemporalArray::Datetime64(Arc::new( + ArrowType::Time32(unit) => Array::TemporalArray(TemporalArray::Datetime32(Arc::new( DatetimeArray { - data: decode_datetime64_plain(buf)?.into(), + data: decode_datetime32_plain(buf)?.into(), null_mask: mask, - time_unit: Default::default(), + time_unit: unit.clone(), }, ))), + #[cfg(feature = "datetime")] + ArrowType::Timestamp(unit, _) | ArrowType::Time64(unit) => { + Array::TemporalArray(TemporalArray::Datetime64(Arc::new(DatetimeArray { + data: decode_datetime64_plain(buf)?.into(), + null_mask: mask, + time_unit: unit.clone(), + }))) + } // decimals #[cfg(feature = "decimal")] ArrowType::Decimal32(precision, scale) => { - let data = decode_int32_plain(buf)?; + let data = match (physical, layout) { + (ParquetPhysicalType::Int32, _) => decode_int32_plain(buf)?, + (ParquetPhysicalType::FixedLenByteArray, ValueLayout::Fixed(width)) => { + decode_decimal_fixed_plain::(buf, width)? + } + _ => return Err(decimal_storage_error(physical)), + }; Array::NumericArray(NumericArray::Decimal32(Arc::new(DecimalArray { data: data.into(), null_mask: mask, @@ -753,7 +713,13 @@ fn decode_column( } #[cfg(feature = "decimal")] ArrowType::Decimal64(precision, scale) => { - let data = decode_int64_plain(buf)?; + let data = match (physical, layout) { + (ParquetPhysicalType::Int64, _) => decode_int64_plain(buf)?, + (ParquetPhysicalType::FixedLenByteArray, ValueLayout::Fixed(width)) => { + decode_decimal_fixed_plain::(buf, width)? + } + _ => return Err(decimal_storage_error(physical)), + }; Array::NumericArray(NumericArray::Decimal64(Arc::new(DecimalArray { data: data.into(), null_mask: mask, @@ -763,7 +729,12 @@ fn decode_column( } #[cfg(feature = "decimal")] ArrowType::Decimal128(precision, scale) => { - let data = decode_decimal128_plain(buf)?; + let data = match (physical, layout) { + (ParquetPhysicalType::FixedLenByteArray, ValueLayout::Fixed(width)) => { + decode_decimal_fixed_plain::(buf, width)? + } + _ => return Err(decimal_storage_error(physical)), + }; Array::NumericArray(NumericArray::Decimal128(Arc::new(DecimalArray { data: data.into(), null_mask: mask, @@ -778,53 +749,6 @@ fn decode_column( }) } -// categorical builders - -#[cfg(any( - not(feature = "default_categorical_8"), - feature = "extended_categorical" -))] -fn build_cat32(idx: Vec64, dict_raw: &[Vec], mask: Option) -> Array { - let dict = dict_raw - .iter() - .map(|b| String::from_utf8(b.clone()).unwrap()) - .collect::>() - .into(); - Array::TextArray(TextArray::Categorical32(Arc::new(CategoricalArray { - data: idx.into(), - unique_values: dict, - null_mask: mask, - }))) -} - -#[cfg(feature = "default_categorical_8")] -fn build_cat8(idx: Vec64, dict_raw: &[Vec], mask: Option) -> Array { - let dict = dict_raw - .iter() - .map(|b| String::from_utf8(b.clone()).unwrap()) - .collect::>(); - let idx8: Vec64 = idx.iter().map(|&v| v as u8).collect(); - Array::TextArray(TextArray::Categorical8(Arc::new(CategoricalArray { - data: idx8.into(), - unique_values: dict, - null_mask: mask, - }))) -} - -#[cfg(all(feature = "extended_categorical", feature = "large_string"))] -fn build_cat64(idx: Vec64, dict_raw: &[Vec], mask: Option) -> Array { - let dict = dict_raw - .iter() - .map(|b| String::from_utf8(b.clone()).unwrap()) - .collect::>() - .into(); - Array::TextArray(TextArray::Categorical64(Arc::new(CategoricalArray { - data: idx.into(), - unique_values: dict, - null_mask: mask, - }))) -} - // RLE/bit-packed Hybrid decoder fn decode_hybrid(buf: &[u8], bit_width: u8, n: usize) -> Result, IoError> { @@ -896,7 +820,6 @@ fn read_uleb128(buf: &[u8]) -> Result<(u64, usize), IoError> { Err(IoError::Format("ULEB128 overflow/truncate".into())) } - // Misc helpers /// Resolve the column chunk's codec id. Uncompressed is `None`. Codecs @@ -915,7 +838,6 @@ fn map_codec(id: i32) -> Result, IoError> { } } - /// Split a PLAIN BYTE_ARRAY dictionary page body into its entries. fn parse_dictionary_values(buf: &[u8]) -> Result>, IoError> { let mut c = Cursor::new(buf); @@ -1012,6 +934,7 @@ fn parse_schema_element(r: &mut R) -> Result { let mut scale = None; let mut field_id = None; let mut num_children = None; + let mut logical_type = None; loop { let (tpe, id) = thrift_read_field_begin(r, &mut last)?; @@ -1020,10 +943,14 @@ fn parse_schema_element(r: &mut R) -> Result { } match id { 1 => { - type_ = Some( - ParquetPhysicalType::from_i32(thrift_read_i32(r)?) - .ok_or_else(|| IoError::Format("Invalid type_".into()))?, - ) + let v = thrift_read_i32(r)?; + type_ = Some(ParquetPhysicalType::from_i32(v).ok_or_else(|| { + if v == 3 { + IoError::UnsupportedType("INT96 columns are not supported".into()) + } else { + IoError::Format(format!("invalid physical type {v}")) + } + })?) } 2 => type_length = Some(thrift_read_i32(r)?), 3 => repetition_type = Some(thrift_read_i32(r)?), @@ -1033,6 +960,7 @@ fn parse_schema_element(r: &mut R) -> Result { 7 => scale = Some(thrift_read_i32(r)?), 8 => precision = Some(thrift_read_i32(r)?), 9 => field_id = Some(thrift_read_i32(r)?), + 10 => logical_type = parse_logical_type(r)?, _ => thrift_skip_field(r, tpe)?, } } @@ -1042,6 +970,7 @@ fn parse_schema_element(r: &mut R) -> Result { repetition_type: repetition_type.unwrap_or(0), type_, converted_type, + logical_type, type_length, precision, scale, @@ -1050,6 +979,125 @@ fn parse_schema_element(r: &mut R) -> Result { }) } +/// Parse the `LogicalType` union of a schema element. +/// +/// The union holds one struct field whose id names the annotation: STRING +/// 1, DECIMAL 5, DATE 6, TIME 7, TIMESTAMP 8, INTEGER 10. TIME and +/// TIMESTAMP carry `isAdjustedToUTC` and a `TimeUnit` union of MILLIS 1, +/// MICROS 2 and NANOS 3. Annotations outside the supported set, such as +/// MAP, LIST, JSON and UUID, return `None` and the column falls back to +/// its converted type. +fn parse_logical_type(r: &mut R) -> Result, IoError> { + let mut last = 0i16; + let mut logical = None; + loop { + let (tpe, id) = thrift_read_field_begin(r, &mut last)?; + if tpe == 0 { + break; + } + if tpe != TC_STRUCT { + thrift_skip_field(r, tpe)?; + continue; + } + logical = match id { + 1 => { + thrift_skip_value(r, TC_STRUCT)?; + Some(ParquetLogicalType::Utf8) + } + #[cfg(feature = "decimal")] + 5 => { + let mut inner = 0i16; + let mut scale = 0i32; + let mut precision = 0i32; + loop { + let (t, f) = thrift_read_field_begin(r, &mut inner)?; + if t == 0 { + break; + } + match f { + 1 => scale = thrift_read_i32(r)?, + 2 => precision = thrift_read_i32(r)?, + _ => thrift_skip_field(r, t)?, + } + } + Some(ParquetLogicalType::Decimal { + precision: precision as u8, + scale: scale as i8, + }) + } + #[cfg(feature = "datetime")] + 6 => { + thrift_skip_value(r, TC_STRUCT)?; + Some(ParquetLogicalType::Date32) + } + #[cfg(feature = "datetime")] + 7 | 8 => { + let mut inner = 0i16; + let mut unit = 0i16; + loop { + let (t, f) = thrift_read_field_begin(r, &mut inner)?; + if t == 0 { + break; + } + match f { + 2 if t == TC_STRUCT => { + let mut unit_last = 0i16; + loop { + let (ut, uf) = thrift_read_field_begin(r, &mut unit_last)?; + if ut == 0 { + break; + } + unit = uf; + thrift_skip_field(r, ut)?; + } + } + _ => thrift_skip_field(r, t)?, + } + } + match (id, unit) { + (7, 1) => Some(ParquetLogicalType::TimeMillis), + (7, 2) => Some(ParquetLogicalType::TimeMicros), + (7, 3) => Some(ParquetLogicalType::TimeNanos), + (8, 1) => Some(ParquetLogicalType::TimestampMillis), + (8, 2) => Some(ParquetLogicalType::TimestampMicros), + (8, 3) => Some(ParquetLogicalType::TimestampNanos), + _ => None, + } + } + 10 => { + let mut inner = 0i16; + let mut bit_width = 0u8; + let mut is_signed = true; + loop { + let (t, f) = thrift_read_field_begin(r, &mut inner)?; + if t == 0 { + break; + } + match (f, t) { + (1, TC_BYTE) => { + let mut b = [0u8; 1]; + r.read_exact(&mut b)?; + bit_width = b[0]; + } + (2, TC_BOOL_TRUE) => is_signed = true, + (2, TC_BOOL_FALSE) => is_signed = false, + _ => thrift_skip_field(r, t)?, + } + } + Some(ParquetLogicalType::IntType { + bit_width, + is_signed, + }) + } + _ => { + thrift_skip_value(r, TC_STRUCT)?; + None + } + }; + } + Ok(logical) +} + fn parse_row_group(r: &mut R) -> Result { let mut last = 0i16; let mut columns = Vec::new(); @@ -1550,11 +1598,6 @@ mod tests { use super::*; - /// Build a Vec string dictionary from &strs. - fn dict(strings: &[&str]) -> Vec> { - strings.iter().map(|s| s.as_bytes().to_vec()).collect() - } - #[test] fn hybrid_rle_run() { // pattern: 6× value 3, bit-width = 2 @@ -1585,61 +1628,6 @@ mod tests { let out = super::decode_hybrid(buf, bit_width, expect.len()).unwrap(); assert_eq!(out.as_slice(), expect.as_slice()); } - #[cfg(not(feature = "default_categorical_8"))] - #[test] - fn decode_column_categorical_rle_dictionary() { - let dict_raw = dict(&["foo", "bar"]); - let idx: Vec = vec![0, 1, 1, 0]; - let encoded: Vec = idx.iter().flat_map(|v| v.to_le_bytes()).collect(); - - let def_levels = vec![true; idx.len()]; - - let array = super::decode_column( - &ArrowType::Dictionary(CategoricalIndexType::UInt32), - &dict_raw, - &encoded, - idx.len(), - def_levels, - ) - .expect("decode_column failed"); - - match array { - Array::TextArray(TextArray::Categorical32(cat)) => { - assert_eq!(cat.data.as_slice(), idx.as_slice()); - let uniq: Vec<_> = cat.unique_values.iter().collect(); - assert_eq!(uniq, vec!["foo", "bar"]); - } - _ => panic!("unexpected array variant {:?}", array), - } - } - - #[cfg(feature = "default_categorical_8")] - #[test] - fn decode_column_categorical_rle_dictionary() { - let dict_raw = dict(&["foo", "bar"]); - let idx: Vec = vec![0, 1, 1, 0]; - let encoded: Vec = idx.iter().flat_map(|v| v.to_le_bytes()).collect(); - - let def_levels = vec![true; idx.len()]; - - let array = super::decode_column( - &ArrowType::Dictionary(CategoricalIndexType::UInt8), - &dict_raw, - &encoded, - idx.len(), - def_levels, - ) - .expect("decode_column failed"); - - match array { - Array::TextArray(TextArray::Categorical8(cat)) => { - assert_eq!(cat.data.as_slice(), &[0u8, 1, 1, 0]); - let uniq: Vec<_> = cat.unique_values.iter().collect(); - assert_eq!(uniq, vec!["foo", "bar"]); - } - _ => panic!("unexpected array variant {:?}", array), - } - } #[test] fn decode_column_plain_int32() { @@ -1651,7 +1639,14 @@ mod tests { } let def_levels = vec![true; values.len()]; - let array = decode_column(&ArrowType::Int32, &[], &buf, values.len(), def_levels.clone()) + let array = decode_column( + &ArrowType::Int32, + ParquetPhysicalType::Int32, + ValueLayout::Fixed(4), + &buf, + values.len(), + def_levels.clone(), + ) .unwrap(); match array { @@ -1669,7 +1664,14 @@ mod tests { // one byte per row, as the page readers unpack boolean pages let bytes: Vec = bits.iter().map(|&b| b as u8).collect(); let def_levels = vec![true; bits.len()]; - let array = decode_column(&ArrowType::Boolean, &[], &bytes, bits.len(), def_levels) + let array = decode_column( + &ArrowType::Boolean, + ParquetPhysicalType::Boolean, + ValueLayout::Fixed(1), + &bytes, + bits.len(), + def_levels, + ) .unwrap(); match array { diff --git a/rust/src/models/types/parquet.rs b/rust/src/models/types/parquet.rs index b01e538..0f8be28 100644 --- a/rust/src/models/types/parquet.rs +++ b/rust/src/models/types/parquet.rs @@ -13,7 +13,7 @@ use crate::error::IoError; #[cfg(feature = "datetime")] use minarrow::TimeUnit; -use minarrow::{ArrowType, ffi::arrow_dtype::CategoricalIndexType}; +use minarrow::ArrowType; /// Parquet physical types as defined in `parquet.thrift`. /// @@ -70,12 +70,9 @@ pub(crate) enum ParquetLogicalType { NoneType, /// UTF-8 encoded string. Utf8, - /// 32-bit date - days since epoch + /// DATE - days since the Unix epoch, stored as INT32. #[cfg(feature = "datetime")] Date32, - /// 64-bit date - milliseconds since epoch - #[cfg(feature = "datetime")] - Date64, /// 64-bit timestamp - milliseconds since epoch #[cfg(feature = "datetime")] TimestampMillis, @@ -298,17 +295,9 @@ pub(crate) fn arrow_type_to_parquet( is_signed: false, }, )), - #[cfg(any( - not(feature = "default_categorical_8"), - feature = "extended_categorical" - ))] - ArrowType::Dictionary(CategoricalIndexType::UInt32) => { - Ok((ParquetPhysicalType::ByteArray, ParquetLogicalType::Utf8)) - } - #[cfg(feature = "default_categorical_8")] - ArrowType::Dictionary(CategoricalIndexType::UInt8) => { - Ok((ParquetPhysicalType::ByteArray, ParquetLogicalType::Utf8)) - } + // Parquet has no dictionary type. A categorical column is a UTF8 + // string column whose pages the writer dictionary-encodes. + ArrowType::Dictionary(_) => Ok((ParquetPhysicalType::ByteArray, ParquetLogicalType::Utf8)), #[cfg(feature = "decimal")] ArrowType::Decimal32(p, s) => Ok(( ParquetPhysicalType::Int32, @@ -332,11 +321,15 @@ pub(crate) fn arrow_type_to_parquet( ArrowType::Utf8View => Ok((ParquetPhysicalType::ByteArray, ParquetLogicalType::Utf8)), #[cfg(feature = "datetime")] ArrowType::Date32 => Ok((ParquetPhysicalType::Int32, ParquetLogicalType::Date32)), + // Parquet DATE is a 32-bit day count, so Date64 milliseconds are + // carried into days on write. See `temporal_unit_scale`. #[cfg(feature = "datetime")] - ArrowType::Date64 => Ok((ParquetPhysicalType::Int64, ParquetLogicalType::Date64)), + ArrowType::Date64 => Ok((ParquetPhysicalType::Int32, ParquetLogicalType::Date32)), #[cfg(feature = "datetime")] ArrowType::Timestamp(unit, _) => match unit { - TimeUnit::Milliseconds => Ok(( + // Parquet has no seconds unit, so seconds are scaled to + // milliseconds on write. + TimeUnit::Seconds | TimeUnit::Milliseconds => Ok(( ParquetPhysicalType::Int64, ParquetLogicalType::TimestampMillis, )), @@ -348,30 +341,16 @@ pub(crate) fn arrow_type_to_parquet( ParquetPhysicalType::Int64, ParquetLogicalType::TimestampNanos, )), - TimeUnit::Seconds => Ok(( - ParquetPhysicalType::Int64, - ParquetLogicalType::TimestampMillis, - )), // best-effort - TimeUnit::Days => Ok((ParquetPhysicalType::Int64, ParquetLogicalType::Date64)), - }, - #[cfg(feature = "datetime")] - ArrowType::Time32(unit) => match unit { - TimeUnit::Milliseconds => { - Ok((ParquetPhysicalType::Int32, ParquetLogicalType::TimeMillis)) - } - TimeUnit::Microseconds => { - Ok((ParquetPhysicalType::Int32, ParquetLogicalType::TimeMicros)) - } - TimeUnit::Nanoseconds => { - Ok((ParquetPhysicalType::Int32, ParquetLogicalType::TimeNanos)) - } - TimeUnit::Seconds => Ok((ParquetPhysicalType::Int32, ParquetLogicalType::TimeMillis)), /* best-effort */ + // A timestamp counted in days is a date. TimeUnit::Days => Ok((ParquetPhysicalType::Int32, ParquetLogicalType::Date32)), }, + // Parquet fixes the storage width by unit: TIME(MILLIS) is INT32, + // TIME(MICROS) and TIME(NANOS) are INT64, whatever width the Arrow + // column uses. Seconds are scaled to milliseconds on write. #[cfg(feature = "datetime")] - ArrowType::Time64(unit) => match unit { - TimeUnit::Milliseconds => { - Ok((ParquetPhysicalType::Int64, ParquetLogicalType::TimeMillis)) + ArrowType::Time32(unit) | ArrowType::Time64(unit) => match unit { + TimeUnit::Seconds | TimeUnit::Milliseconds => { + Ok((ParquetPhysicalType::Int32, ParquetLogicalType::TimeMillis)) } TimeUnit::Microseconds => { Ok((ParquetPhysicalType::Int64, ParquetLogicalType::TimeMicros)) @@ -379,25 +358,39 @@ pub(crate) fn arrow_type_to_parquet( TimeUnit::Nanoseconds => { Ok((ParquetPhysicalType::Int64, ParquetLogicalType::TimeNanos)) } - TimeUnit::Seconds => Ok((ParquetPhysicalType::Int64, ParquetLogicalType::TimeMillis)), /* best-effort */ - TimeUnit::Days => Ok((ParquetPhysicalType::Int64, ParquetLogicalType::Date64)), + TimeUnit::Days => Err(IoError::UnsupportedType( + "time of day counted in days has no Parquet type".into(), + )), }, ArrowType::Null => Err(IoError::UnsupportedType( "Null type is not supported".into(), )), #[cfg(feature = "datetime")] - ArrowType::Duration32(_) => panic!("Duration does not map to a parquet type."), + ArrowType::Duration32(_) | ArrowType::Duration64(_) => Err(IoError::UnsupportedType( + "Duration has no Parquet logical type".into(), + )), #[cfg(feature = "datetime")] - ArrowType::Duration64(_) => panic!("Duration does not map to a parquet type."), + ArrowType::Interval(_) => Err(IoError::UnsupportedType( + "Interval is not supported".into(), + )), + } +} + +/// Multiplier and divisor that carry an Arrow temporal value into the unit +/// of its Parquet annotation from [`arrow_type_to_parquet`]. +/// +/// Most units match and scale by one. Seconds become milliseconds because +/// Parquet has no seconds unit, and Date64 milliseconds become the day +/// count that Parquet DATE stores. +pub(crate) fn temporal_unit_scale(ty: &ArrowType) -> (i64, i64) { + match ty { + #[cfg(feature = "datetime")] + ArrowType::Date64 => (1, 86_400_000), #[cfg(feature = "datetime")] - ArrowType::Interval(_) => panic!("Interval does not map to a parquet type."), - #[cfg(all(feature = "extended_categorical", feature = "extended_numeric_types"))] - &minarrow::ArrowType::Dictionary( - minarrow::ffi::arrow_dtype::CategoricalIndexType::UInt16, - ) - | &minarrow::ArrowType::Dictionary( - minarrow::ffi::arrow_dtype::CategoricalIndexType::UInt64, - ) => panic!(), + ArrowType::Timestamp(TimeUnit::Seconds, _) + | ArrowType::Time32(TimeUnit::Seconds) + | ArrowType::Time64(TimeUnit::Seconds) => (1000, 1), + _ => (1, 1), } } @@ -480,8 +473,6 @@ pub(crate) fn parquet_to_arrow_type( #[cfg(feature = "datetime")] (ParquetPhysicalType::Int32, Some(ParquetLogicalType::Date32)) => Ok(ArrowType::Date32), #[cfg(feature = "datetime")] - (ParquetPhysicalType::Int64, Some(ParquetLogicalType::Date64)) => Ok(ArrowType::Date64), - #[cfg(feature = "datetime")] (ParquetPhysicalType::Int64, Some(ParquetLogicalType::TimestampMillis)) => { Ok(ArrowType::Timestamp(TimeUnit::Milliseconds, None)) } @@ -529,23 +520,25 @@ pub(crate) fn parquet_to_arrow_type( ParquetPhysicalType::Int64, Some(ParquetLogicalType::Decimal { precision, scale }), ) => Ok(ArrowType::Decimal64(precision, scale)), + // A fixed-length decimal may hold any precision. The narrowest + // Arrow decimal that fits the precision is used, matching the + // INT32 and INT64 forms above. #[cfg(feature = "decimal")] ( ParquetPhysicalType::FixedLenByteArray, Some(ParquetLogicalType::Decimal { precision, scale }), - ) => Ok(ArrowType::Decimal128(precision, scale)), + ) => Ok(match precision { + 0..=9 => ArrowType::Decimal32(precision, scale), + 10..=18 => ArrowType::Decimal64(precision, scale), + _ => ArrowType::Decimal128(precision, scale), + }), // Floats (ParquetPhysicalType::Float, _) => Ok(ArrowType::Float32), (ParquetPhysicalType::Double, _) => Ok(ArrowType::Float64), // Strings - always logical UTF8/Utf8 - #[cfg(not(feature = "large_string"))] (ParquetPhysicalType::ByteArray, Some(ParquetLogicalType::Utf8)) => Ok(ArrowType::String), - #[cfg(feature = "large_string")] - (ParquetPhysicalType::ByteArray, Some(ParquetLogicalType::Utf8)) => { - Ok(ArrowType::LargeString) - } // Fallback -- treat byte array without logical utf8 as unsupported (ParquetPhysicalType::ByteArray, None) => { diff --git a/rust/src/models/writers/parquet.rs b/rust/src/models/writers/parquet.rs index 05d243d..191fa70 100644 --- a/rust/src/models/writers/parquet.rs +++ b/rust/src/models/writers/parquet.rs @@ -46,6 +46,8 @@ use std::io::{Seek, Write}; #[cfg(feature = "datetime")] use minarrow::TemporalArray; +#[cfg(feature = "datetime")] +use minarrow::ArrowType; use minarrow::{Array, NumericArray, Table, TextArray}; use crate::compression::{Compression, compress}; @@ -55,15 +57,17 @@ use crate::error::IoError; use crate::models::encoders::parquet::data::encode_large_string_plain; use crate::models::encoders::parquet::data::{ encode_bool_bitpacked, encode_float32_plain, encode_float64_plain, encode_int32_plain, - encode_int64_plain, encode_string_plain, encode_uint32_as_int32_plain, - encode_uint64_as_int64_plain, + encode_int64_plain, encode_string_plain, encode_temporal_plain, + encode_uint32_as_int32_plain, encode_uint64_as_int64_plain, }; use crate::models::encoders::parquet::metadata::{ ColumnChunkMeta, ColumnMetadata, DataPageHeaderV2, DictionaryPageHeader, FileMetaData, PageHeader, PageType, RowGroupMeta, SchemaElement, Statistics, }; use crate::models::types::parquet::ParquetLogicalType::{self}; -use crate::models::types::parquet::{ParquetEncoding, ParquetPhysicalType, arrow_type_to_parquet}; +use crate::models::types::parquet::{ + ParquetEncoding, ParquetPhysicalType, arrow_type_to_parquet, temporal_unit_scale, +}; // Chunk size for page splitting pub const PARQUET_PAGE_CHUNK_SIZE: usize = 32_768; @@ -110,6 +114,7 @@ pub fn write_parquet_table( scale: None, field_id: None, num_children: Some(table.cols.len() as i32), + logical_type: None, }); for (i, c) in table.cols.iter().enumerate() { let (physical, logical) = arrow_type_to_parquet(&c.field.dtype).unwrap(); @@ -119,6 +124,7 @@ pub fn write_parquet_table( repetition_type: if c.field.nullable { 1 } else { 0 }, // OPTIONAL / REQUIRED type_: Some(physical), converted_type: logical_to_converted(&logical), + logical_type: Some(logical), type_length, precision, scale, @@ -137,6 +143,27 @@ pub fn write_parquet_table( // Column loop, multi-page support for col in &table.cols { + // Temporal columns are stored in the unit and width of their Parquet + // annotation, so their values may need scaling or narrowing on the + // way out. Every other type is stored as-is. + let (phys, _) = arrow_type_to_parquet(&col.field.dtype)?; + let temporal = match &col.field.dtype { + #[cfg(feature = "datetime")] + ArrowType::Date32 + | ArrowType::Date64 + | ArrowType::Timestamp(_, _) + | ArrowType::Time32(_) + | ArrowType::Time64(_) => true, + _ => false, + }; + let (unit_mul, unit_div) = temporal_unit_scale(&col.field.dtype); + let encode_temporal32 = |data: &[i32], out: &mut Vec| { + encode_temporal_plain(data, unit_mul, unit_div, phys, out) + }; + let encode_temporal64 = |data: &[i64], out: &mut Vec| { + encode_temporal_plain(data, unit_mul, unit_div, phys, out) + }; + let mut dictionary_page_offset = None; let mut encodings = vec![ParquetEncoding::Plain]; @@ -224,12 +251,18 @@ pub fn write_parquet_table( // encode the raw values for this slice match &col.array { Array::NumericArray(n) => match n { + NumericArray::Int32(a) if temporal => { + encode_valid!(encode_temporal32, &a.data[start..end])? + } NumericArray::Int32(a) => { encode_valid!(encode_int32_plain, &a.data[start..end]) } NumericArray::UInt32(a) => { encode_valid!(encode_uint32_as_int32_plain, &a.data[start..end]) } + NumericArray::Int64(a) if temporal => { + encode_valid!(encode_temporal64, &a.data[start..end])? + } NumericArray::Int64(a) => { encode_valid!(encode_int64_plain, &a.data[start..end]) } @@ -301,15 +334,11 @@ pub fn write_parquet_table( )?, #[cfg(feature = "datetime")] Array::TemporalArray(TemporalArray::Datetime32(a)) => { - use crate::models::encoders::parquet::data::encode_datetime32_plain; - - encode_valid!(encode_datetime32_plain, &a.data[start..end]) + encode_valid!(encode_temporal32, &a.data[start..end])? } #[cfg(feature = "datetime")] Array::TemporalArray(TemporalArray::Datetime64(a)) => { - use crate::models::encoders::parquet::data::encode_datetime64_plain; - - encode_valid!(encode_datetime64_plain, &a.data[start..end]) + encode_valid!(encode_temporal64, &a.data[start..end])? } #[cfg(any( not(feature = "default_categorical_8"), @@ -422,7 +451,6 @@ pub fn write_parquet_table( // column-chunk metadata let first_data = recorded_data_page_offset.expect("at least one data page must be emitted"); - let (phys, _) = arrow_type_to_parquet(&col.field.dtype)?; columns_meta.push(ColumnChunkMeta { file_offset: first_data, meta_data: ColumnMetadata { @@ -632,8 +660,6 @@ fn logical_to_converted(log: &ParquetLogicalType) -> Option { #[cfg(feature = "datetime")] ParquetLogicalType::Date32 => 6, #[cfg(feature = "datetime")] - ParquetLogicalType::Date64 => return None, - #[cfg(feature = "datetime")] ParquetLogicalType::TimestampMillis => 9, #[cfg(feature = "datetime")] ParquetLogicalType::TimestampMicros => 10, diff --git a/rust/tests/parquet_nullable_roundtrip.rs b/rust/tests/parquet_nullable_roundtrip.rs index 16bc8ec..1dd935a 100644 --- a/rust/tests/parquet_nullable_roundtrip.rs +++ b/rust/tests/parquet_nullable_roundtrip.rs @@ -10,6 +10,11 @@ //! every page boundary, so the value sections hold non-null entries only //! and the reader has to scatter them back against the definition levels //! page by page. +//! +//! The types that come back are the ones the Parquet schema names, which +//! is not always the Arrow type that went in: categoricals are UTF8 +//! strings, Date64 is a DATE day count, seconds become milliseconds, and +//! TIME columns take the storage width their unit requires. #[cfg(feature = "parquet")] mod parquet_nullable_roundtrip_tests { @@ -24,7 +29,9 @@ mod parquet_nullable_roundtrip_tests { IntegerArray, MaskedArray, NumericArray, StringArray, Table, TextArray, Vec64, }; #[cfg(feature = "datetime")] - use minarrow::{DatetimeArray, TemporalArray}; + use minarrow::{DatetimeArray, TemporalArray, TimeUnit}; + #[cfg(feature = "decimal")] + use minarrow::DecimalArray; /// Rows per column: two full pages plus a partial third page. const N_ROWS: usize = 2 * PARQUET_PAGE_CHUNK_SIZE + 13; @@ -83,11 +90,63 @@ mod parquet_nullable_roundtrip_tests { fn date32_value(i: usize) -> i32 { i as i32 * 3 } + /// Milliseconds with a sub-day remainder, so the DATE day count is the + /// truncated quotient. + #[cfg(feature = "datetime")] + fn date64_value(i: usize) -> i64 { + i as i64 * 86_400_000 + 12_345 + } + #[cfg(feature = "datetime")] + fn timestamp_value(i: usize) -> i64 { + 1_700_000_000_000 + i as i64 * 1_001 + } + #[cfg(feature = "datetime")] + fn time_value(i: usize) -> i64 { + (i % 86_400) as i64 * 7 + } + #[cfg(feature = "decimal")] + fn decimal32_value(i: usize) -> i32 { + i as i32 * 125 - 50_000 + } + #[cfg(feature = "decimal")] + fn decimal64_value(i: usize) -> i64 { + i as i64 * 1_000_003 - 7 + } + #[cfg(feature = "decimal")] + fn decimal128_value(i: usize) -> i128 { + (i as i128 - 5) * 1_000_000_000_000_000_000_000 + } fn column(name: &str, dtype: ArrowType, array: Array) -> FieldArray { FieldArray::new(Field::new(name, dtype, true, None), array) } + #[cfg(feature = "datetime")] + fn temporal32(name: &str, dtype: ArrowType, unit: TimeUnit, value: fn(usize) -> i32) -> FieldArray { + column( + name, + dtype, + Array::from_datetime_i32(DatetimeArray::from_vec64( + (0..N_ROWS).map(value).collect::>(), + Some(null_mask()), + Some(unit), + )), + ) + } + + #[cfg(feature = "datetime")] + fn temporal64(name: &str, dtype: ArrowType, unit: TimeUnit, value: fn(usize) -> i64) -> FieldArray { + column( + name, + dtype, + Array::from_datetime_i64(DatetimeArray::from_vec64( + (0..N_ROWS).map(value).collect::>(), + Some(null_mask()), + Some(unit), + )), + ) + } + fn all_types_table() -> Table { let mask = null_mask(); let strings: Vec = (0..N_ROWS).map(string_value).collect(); @@ -181,15 +240,88 @@ mod parquet_nullable_roundtrip_tests { #[cfg(feature = "datetime")] { - // Date64 is left out: Parquet has no 64-bit DATE annotation, so - // the type mapping stores it as a plain INT64. + cols.push(temporal32("date32", ArrowType::Date32, TimeUnit::Days, date32_value)); + cols.push(temporal64("date64", ArrowType::Date64, TimeUnit::Milliseconds, date64_value)); + cols.push(temporal64( + "ts_s", + ArrowType::Timestamp(TimeUnit::Seconds, None), + TimeUnit::Seconds, + timestamp_value, + )); + cols.push(temporal64( + "ts_ms", + ArrowType::Timestamp(TimeUnit::Milliseconds, None), + TimeUnit::Milliseconds, + timestamp_value, + )); + cols.push(temporal64( + "ts_us", + ArrowType::Timestamp(TimeUnit::Microseconds, None), + TimeUnit::Microseconds, + timestamp_value, + )); + cols.push(temporal64( + "ts_ns", + ArrowType::Timestamp(TimeUnit::Nanoseconds, None), + TimeUnit::Nanoseconds, + timestamp_value, + )); + cols.push(temporal32( + "time32_ms", + ArrowType::Time32(TimeUnit::Milliseconds), + TimeUnit::Milliseconds, + |i| time_value(i) as i32, + )); + cols.push(temporal32( + "time32_us", + ArrowType::Time32(TimeUnit::Microseconds), + TimeUnit::Microseconds, + |i| time_value(i) as i32, + )); + cols.push(temporal64( + "time64_ms", + ArrowType::Time64(TimeUnit::Milliseconds), + TimeUnit::Milliseconds, + time_value, + )); + cols.push(temporal64( + "time64_ns", + ArrowType::Time64(TimeUnit::Nanoseconds), + TimeUnit::Nanoseconds, + time_value, + )); + } + + #[cfg(feature = "decimal")] + { + cols.push(column( + "dec32", + ArrowType::Decimal32(7, 2), + Array::from_decimal32(DecimalArray::from_vec64( + (0..N_ROWS).map(decimal32_value).collect::>(), + Some(mask.clone()), + 7, + 2, + )), + )); + cols.push(column( + "dec64", + ArrowType::Decimal64(18, 4), + Array::from_decimal64(DecimalArray::from_vec64( + (0..N_ROWS).map(decimal64_value).collect::>(), + Some(mask.clone()), + 18, + 4, + )), + )); cols.push(column( - "date32", - ArrowType::Date32, - Array::from_datetime_i32(DatetimeArray::from_vec64( - (0..N_ROWS).map(date32_value).collect::>(), + "dec128", + ArrowType::Decimal128(38, 6), + Array::from_decimal128(DecimalArray::from_vec64( + (0..N_ROWS).map(decimal128_value).collect::>(), Some(mask.clone()), - None, + 38, + 6, )), )); } @@ -212,6 +344,30 @@ mod parquet_nullable_roundtrip_tests { .unwrap_or_else(|| panic!("column {name} missing")) } + #[cfg(feature = "datetime")] + fn assert_temporal32(out: &Table, name: &str, dtype: ArrowType, unit: TimeUnit, value: impl Fn(usize) -> i32) { + assert_eq!(col(out, name).field.dtype, dtype, "{name} dtype"); + match &col(out, name).array { + Array::TemporalArray(TemporalArray::Datetime32(a)) => { + assert_eq!(a.time_unit, unit, "{name} unit"); + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(value), "{name}"); + } + other => panic!("{name}: {other:?}"), + } + } + + #[cfg(feature = "datetime")] + fn assert_temporal64(out: &Table, name: &str, dtype: ArrowType, unit: TimeUnit, value: impl Fn(usize) -> i64) { + assert_eq!(col(out, name).field.dtype, dtype, "{name} dtype"); + match &col(out, name).array { + Array::TemporalArray(TemporalArray::Datetime64(a)) => { + assert_eq!(a.time_unit, unit, "{name} unit"); + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(value), "{name}"); + } + other => panic!("{name}: {other:?}"), + } + } + fn assert_all_types(out: &Table) { let expected_nulls = (0..N_ROWS).filter(|&i| !is_valid(i)).count(); assert_eq!(out.n_rows, N_ROWS); @@ -263,40 +419,119 @@ mod parquet_nullable_roundtrip_tests { } other => panic!("bool: {other:?}"), } - let strings: Vec> = match &col(out, "string").array { - Array::TextArray(TextArray::String32(a)) => { - (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect() - } - #[cfg(feature = "large_string")] - Array::TextArray(TextArray::String64(a)) => { - (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect() - } + assert_eq!(col(out, "string").field.dtype, ArrowType::String); + match &col(out, "string").array { + Array::TextArray(TextArray::String32(a)) => assert_eq!( + (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect::>(), + expected(string_value) + ), other => panic!("string: {other:?}"), - }; - assert_eq!(strings, expected(string_value)); + } - let categories: Vec> = match &col(out, "category").array { - #[cfg(feature = "default_categorical_8")] - Array::TextArray(TextArray::Categorical8(a)) => { - (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect() - } - #[cfg(not(feature = "default_categorical_8"))] - Array::TextArray(TextArray::Categorical32(a)) => { - (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect() - } + // A categorical column is a UTF8 string column in Parquet, so it + // reads back as one. + assert_eq!(col(out, "category").field.dtype, ArrowType::String); + match &col(out, "category").array { + Array::TextArray(TextArray::String32(a)) => assert_eq!( + (0..N_ROWS).map(|i| a.get(i).map(str::to_owned)).collect::>(), + expected(|i| category_value(i).to_owned()) + ), other => panic!("category: {other:?}"), - }; - assert_eq!(categories, expected(|i| category_value(i).to_owned())); + } #[cfg(feature = "datetime")] { - assert_eq!(col(out, "date32").field.dtype, ArrowType::Date32); - match &col(out, "date32").array { - Array::TemporalArray(TemporalArray::Datetime32(a)) => assert_eq!( - (0..N_ROWS).map(|i| a.get(i)).collect::>(), - expected(date32_value) - ), - other => panic!("date32: {other:?}"), + assert_temporal32(out, "date32", ArrowType::Date32, TimeUnit::Days, date32_value); + // DATE stores days, so Date64 milliseconds come back as Date32. + assert_temporal32(out, "date64", ArrowType::Date32, TimeUnit::Days, |i| { + (date64_value(i) / 86_400_000) as i32 + }); + // Parquet has no seconds unit, so seconds come back as milliseconds. + assert_temporal64( + out, + "ts_s", + ArrowType::Timestamp(TimeUnit::Milliseconds, None), + TimeUnit::Milliseconds, + |i| timestamp_value(i) * 1000, + ); + assert_temporal64( + out, + "ts_ms", + ArrowType::Timestamp(TimeUnit::Milliseconds, None), + TimeUnit::Milliseconds, + timestamp_value, + ); + assert_temporal64( + out, + "ts_us", + ArrowType::Timestamp(TimeUnit::Microseconds, None), + TimeUnit::Microseconds, + timestamp_value, + ); + assert_temporal64( + out, + "ts_ns", + ArrowType::Timestamp(TimeUnit::Nanoseconds, None), + TimeUnit::Nanoseconds, + timestamp_value, + ); + assert_temporal32( + out, + "time32_ms", + ArrowType::Time32(TimeUnit::Milliseconds), + TimeUnit::Milliseconds, + |i| time_value(i) as i32, + ); + // TIME(MICROS) is INT64 in Parquet, so Time32 microseconds widen. + assert_temporal64( + out, + "time32_us", + ArrowType::Time64(TimeUnit::Microseconds), + TimeUnit::Microseconds, + time_value, + ); + // TIME(MILLIS) is INT32 in Parquet, so Time64 milliseconds narrow. + assert_temporal32( + out, + "time64_ms", + ArrowType::Time32(TimeUnit::Milliseconds), + TimeUnit::Milliseconds, + |i| time_value(i) as i32, + ); + assert_temporal64( + out, + "time64_ns", + ArrowType::Time64(TimeUnit::Nanoseconds), + TimeUnit::Nanoseconds, + time_value, + ); + } + + #[cfg(feature = "decimal")] + { + assert_eq!(col(out, "dec32").field.dtype, ArrowType::Decimal32(7, 2)); + match &col(out, "dec32").array { + Array::NumericArray(NumericArray::Decimal32(a)) => { + assert_eq!((a.precision, a.scale), (7, 2)); + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(decimal32_value)); + } + other => panic!("dec32: {other:?}"), + } + assert_eq!(col(out, "dec64").field.dtype, ArrowType::Decimal64(18, 4)); + match &col(out, "dec64").array { + Array::NumericArray(NumericArray::Decimal64(a)) => { + assert_eq!((a.precision, a.scale), (18, 4)); + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(decimal64_value)); + } + other => panic!("dec64: {other:?}"), + } + assert_eq!(col(out, "dec128").field.dtype, ArrowType::Decimal128(38, 6)); + match &col(out, "dec128").array { + Array::NumericArray(NumericArray::Decimal128(a)) => { + assert_eq!((a.precision, a.scale), (38, 6)); + assert_eq!((0..N_ROWS).map(|i| a.get(i)).collect::>(), expected(decimal128_value)); + } + other => panic!("dec128: {other:?}"), } } } diff --git a/rust/tests/pyarrow_parquet.rs b/rust/tests/pyarrow_parquet.rs index 4517912..84dde80 100644 --- a/rust/tests/pyarrow_parquet.rs +++ b/rust/tests/pyarrow_parquet.rs @@ -19,12 +19,25 @@ mod pyarrow_parquet_tests { use lightstream::models::readers::parquet::{load_parquet_table, load_parquet_table_cols}; use minarrow::{Array, ArrowType, MaskedArray, NumericArray, Table, TextArray}; + #[cfg(all(feature = "datetime", feature = "decimal", feature = "snappy"))] + use minarrow::{TemporalArray, TimeUnit}; fn open(name: &str) -> BufReader { let path = format!("{}/pyarrow-roundtrip/{name}", env!("CARGO_MANIFEST_DIR")); BufReader::new(File::open(&path).unwrap_or_else(|e| panic!("open {path}: {e}"))) } + fn column_dtype(table: &Table, name: &str) -> ArrowType { + table + .cols + .iter() + .find(|c| c.field.name == name) + .unwrap_or_else(|| panic!("column {name} missing")) + .field + .dtype + .clone() + } + fn column<'a>(table: &'a Table, name: &str) -> &'a Array { let col = table .cols @@ -73,18 +86,12 @@ mod pyarrow_parquet_tests { } } - /// UTF-8 columns map to `String`, or to `LargeString` when the - /// `large_string` feature is on. fn string_values(table: &Table, name: &str) -> Vec> { match column(table, name) { Array::TextArray(TextArray::String32(a)) => { (0..a.len()).map(|i| a.get(i).map(str::to_owned)).collect() } - #[cfg(feature = "large_string")] - Array::TextArray(TextArray::String64(a)) => { - (0..a.len()).map(|i| a.get(i).map(str::to_owned)).collect() - } - other => panic!("{name}: expected a UTF-8 column, got {other:?}"), + other => panic!("{name}: expected String, got {other:?}"), } } @@ -147,10 +154,7 @@ mod pyarrow_parquet_tests { assert_eq!(table.n_rows, 5); assert_eq!(table.cols.len(), 3); assert_eq!(table.cols[0].field.dtype, ArrowType::Int64); - #[cfg(not(feature = "large_string"))] assert_eq!(table.cols[1].field.dtype, ArrowType::String); - #[cfg(feature = "large_string")] - assert_eq!(table.cols[1].field.dtype, ArrowType::LargeString); assert_eq!(table.cols[2].field.dtype, ArrowType::Float64); assert_eq!( @@ -195,6 +199,107 @@ mod pyarrow_parquet_tests { assert_nullable_table(&table); } + /// Temporal and decimal columns written with pyarrow's defaults. The + /// nanosecond timestamp exists only through the `LogicalType` + /// annotation, and pyarrow stores every decimal as a + /// FIXED_LEN_BYTE_ARRAY sized to its precision. + #[cfg(all(feature = "datetime", feature = "decimal", feature = "snappy"))] + #[test] + fn reads_pyarrow_temporal_and_decimal_columns() { + let table = load_parquet_table(open("pyarrow_temporal_decimal.parquet")).expect("read"); + assert_eq!(table.n_rows, 5); + + fn temporal32(table: &Table, name: &str, unit: TimeUnit) -> Vec> { + match column(table, name) { + Array::TemporalArray(TemporalArray::Datetime32(a)) => { + assert_eq!(a.time_unit, unit, "{name} unit"); + (0..a.len()).map(|i| a.get(i)).collect() + } + other => panic!("{name}: expected Datetime32, got {other:?}"), + } + } + fn temporal64(table: &Table, name: &str, unit: TimeUnit) -> Vec> { + match column(table, name) { + Array::TemporalArray(TemporalArray::Datetime64(a)) => { + assert_eq!(a.time_unit, unit, "{name} unit"); + (0..a.len()).map(|i| a.get(i)).collect() + } + other => panic!("{name}: expected Datetime64, got {other:?}"), + } + } + + assert_eq!(column_dtype(&table, "date"), ArrowType::Date32); + assert_eq!( + temporal32(&table, "date", TimeUnit::Days), + [Some(0), Some(1), None, Some(19_000), Some(-5)] + ); + assert_eq!(column_dtype(&table, "time_ms"), ArrowType::Time32(TimeUnit::Milliseconds)); + assert_eq!( + temporal32(&table, "time_ms", TimeUnit::Milliseconds), + [Some(0), Some(1_000), None, Some(43_200_000), Some(86_399_999)] + ); + assert_eq!(column_dtype(&table, "time_us"), ArrowType::Time64(TimeUnit::Microseconds)); + assert_eq!( + temporal64(&table, "time_us", TimeUnit::Microseconds), + [Some(0), None, Some(2_000_000), Some(43_200_000_000), Some(86_399_999_999)] + ); + assert_eq!( + column_dtype(&table, "ts_ms"), + ArrowType::Timestamp(TimeUnit::Milliseconds, None) + ); + assert_eq!( + temporal64(&table, "ts_ms", TimeUnit::Milliseconds), + [Some(0), Some(1_700_000_000_000), None, Some(-1), Some(86_400_000)] + ); + assert_eq!( + column_dtype(&table, "ts_us"), + ArrowType::Timestamp(TimeUnit::Microseconds, None) + ); + assert_eq!( + temporal64(&table, "ts_us", TimeUnit::Microseconds), + [Some(0), Some(1_700_000_000_000_000), None, Some(-1), Some(1)] + ); + assert_eq!( + column_dtype(&table, "ts_ns"), + ArrowType::Timestamp(TimeUnit::Nanoseconds, None) + ); + assert_eq!( + temporal64(&table, "ts_ns", TimeUnit::Nanoseconds), + [Some(0), Some(1_700_000_000_000_000_000), None, Some(-1), Some(1)] + ); + + assert_eq!(column_dtype(&table, "dec32"), ArrowType::Decimal32(7, 2)); + match column(&table, "dec32") { + Array::NumericArray(NumericArray::Decimal32(a)) => assert_eq!( + (0..a.len()).map(|i| a.get(i)).collect::>(), + [Some(125), None, Some(-350), Some(1), Some(9_999_999)] + ), + other => panic!("dec32: {other:?}"), + } + assert_eq!(column_dtype(&table, "dec64"), ArrowType::Decimal64(18, 4)); + match column(&table, "dec64") { + Array::NumericArray(NumericArray::Decimal64(a)) => assert_eq!( + (0..a.len()).map(|i| a.get(i)).collect::>(), + [Some(12_345), Some(-10_000), None, Some(1), Some(123_456_789_012_345_678)] + ), + other => panic!("dec64: {other:?}"), + } + assert_eq!(column_dtype(&table, "dec128"), ArrowType::Decimal128(38, 6)); + match column(&table, "dec128") { + Array::NumericArray(NumericArray::Decimal128(a)) => assert_eq!( + (0..a.len()).map(|i| a.get(i)).collect::>(), + [ + Some(1_000_001), + None, + Some(-1_000_001), + Some(0), + Some(12_345_678_901_234_567_890_123_456_789_012_123_456i128) + ] + ), + other => panic!("dec128: {other:?}"), + } + } + /// Column projection over a pyarrow file selects by name and keeps /// the values of the selected columns intact. #[cfg(feature = "snappy")] From 5ad0a48474466ef52a1f36bd0a4df65e39e1c88c Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Thu, 10 Sep 2026 06:54:03 +0100 Subject: [PATCH 6/7] Update CHANGELOG.md --- CHANGELOG.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b261085..b50a934 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,17 +2,18 @@ Notable changes are recorded from 0.5.0 onward. -## 0.6.2 +## 0.7.0 ### Changed - Minarrow upgraded to 0.18.1 and vec64 0.5.1. - New `decimal` feature forwarding minarrow's `decimal` feature. Decimal32, Decimal64 and Decimal128 columns are supported in Arrow IPC and in Parquet, where they map to the DECIMAL logical type over INT32, INT64 and FIXED_LEN_BYTE_ARRAY. - The Python package pins minarrow and minarrow-pyo3 at 0.18.1. -- The Parquet writer follows the Parquet value layout for nullable columns. Value sections hold non-null values only and DataPageV2 headers count every row in `num_values`, so files with nulls now read in pyarrow and other Parquet readers. Files with nulls written by earlier releases do not read back under this release. -- Parquet column types follow the schema element. Categorical columns are UTF8 string columns with dictionary-encoded pages and read back as strings. Date64 is stored as DATE days, seconds-unit timestamps and times as milliseconds, and TIME columns take the INT32 or INT64 width their unit requires. UTF8 columns read as `String` and widen to `LargeString` only when the data exceeds 32-bit offsets. -- The `LogicalType` schema annotation is written and read, so nanosecond timestamps and times round-trip and pyarrow nanosecond timestamps no longer read as Int64. -- Duration and Interval columns are reported as `UnsupportedType` by the Parquet writer instead of panicking, and INT96 columns are reported by name on read. +- Parquet writer improvements: + - The Parquet writer follows the Parquet value layout for nullable columns. + - Several improvements in categorical, time, and date roundtripping. + - Files with nulls written by earlier releases do not read back under this release. + - Replaced packing for Duration and Interval columns with `UnsupportedType` ### Fixed From c25a04d5afdf9596ddcd6cb7f5d28109c8c4ce3d Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Thu, 10 Sep 2026 06:58:50 +0100 Subject: [PATCH 7/7] Release 0.7.0: carry isAdjustedToUTC on Parquet timestamps and bump the crate and Python package --- python/Cargo.lock | 4 +- python/Cargo.toml | 4 +- python/pyproject.toml | 2 +- rust/Cargo.lock | 2 +- rust/Cargo.toml | 2 +- .../generate_pyarrow_parquet_files.py | 1 + .../pyarrow_temporal_decimal.parquet | Bin 3054 -> 3396 bytes rust/src/models/encoders/parquet/metadata.rs | 27 ++++++--- rust/src/models/readers/parquet.rs | 8 ++- rust/src/models/types/parquet.rs | 52 +++++++++++------- rust/src/models/writers/parquet.rs | 6 +- rust/tests/parquet_nullable_roundtrip.rs | 15 +++++ rust/tests/pyarrow_parquet.rs | 9 +++ 13 files changed, 90 insertions(+), 42 deletions(-) diff --git a/python/Cargo.lock b/python/Cargo.lock index e9cc009..c538bef 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -662,7 +662,7 @@ checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "lightstream" -version = "0.6.2" +version = "0.7.0" dependencies = [ "bytes", "fast-float2", @@ -692,7 +692,7 @@ dependencies = [ [[package]] name = "lightstream-py" -version = "0.6.2" +version = "0.7.0" dependencies = [ "futures-core", "lightstream", diff --git a/python/Cargo.toml b/python/Cargo.toml index 4d36ff2..9aef560 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -2,7 +2,7 @@ cargo-features = ["trim-paths"] [package] name = "lightstream-py" -version = "0.6.2" +version = "0.7.0" edition = "2024" authors = ["Peter G. Bower"] license = "MPL-2.0" @@ -19,7 +19,7 @@ name = "lightstream_py" crate-type = ["cdylib", "rlib"] [dependencies] -lightstream = { version = "0.6", path = "../rust", features = ["csv", "datetime", "decimal", "extended_categorical", "extended_numeric_types", "json", "mmap", "http", "parquet", "protocol", "quic", "snappy", "stdio", "tcp", "tls", "uds", "webtransport", "websocket", "zstd"] } +lightstream = { version = "0.7", path = "../rust", features = ["csv", "datetime", "decimal", "extended_categorical", "extended_numeric_types", "json", "mmap", "http", "parquet", "protocol", "quic", "snappy", "stdio", "tcp", "tls", "uds", "webtransport", "websocket", "zstd"] } # The categorical and numeric feature set mirrors the minarrow-py build. # minarrow-pyo3's dictionary-index conversion needs the extended features, # and they flow through lightstream's flags so its match arms gate in step diff --git a/python/pyproject.toml b/python/pyproject.toml index fc5c59a..5c6f7a1 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "lightstream-io" -version = "0.6.2" +version = "0.7.0" description = "Streaming Arrow I/O for Python - files, sockets, and network transports with zero-copy minarrow interop." readme = "README.md" requires-python = ">=3.9" diff --git a/rust/Cargo.lock b/rust/Cargo.lock index caccf56..9387c3a 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1924,7 +1924,7 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "lightstream" -version = "0.6.2" +version = "0.7.0" dependencies = [ "arrow", "arrow-flight", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index be37ee2..7bf99ed 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lightstream" -version = "0.6.2" +version = "0.7.0" edition = "2024" license = "MPL-2.0" keywords = [ diff --git a/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py b/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py index 00d5178..e30de12 100644 --- a/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py +++ b/rust/pyarrow-roundtrip/generate_pyarrow_parquet_files.py @@ -68,6 +68,7 @@ def temporal_decimal_table() -> pa.Table: "ts_ms": pa.array([0, 1_700_000_000_000, None, -1, 86_400_000], type=pa.timestamp("ms")), "ts_us": pa.array([0, 1_700_000_000_000_000, None, -1, 1], type=pa.timestamp("us")), "ts_ns": pa.array([0, 1_700_000_000_000_000_000, None, -1, 1], type=pa.timestamp("ns")), + "ts_us_utc": pa.array([0, 1_700_000_000_000_000, None, -1, 1], type=pa.timestamp("us", tz="UTC")), "dec32": pa.array( [Decimal("1.25"), None, Decimal("-3.50"), Decimal("0.01"), Decimal("99999.99")], type=pa.decimal32(7, 2), diff --git a/rust/pyarrow-roundtrip/pyarrow_temporal_decimal.parquet b/rust/pyarrow-roundtrip/pyarrow_temporal_decimal.parquet index 5a8bef6c49ff2827c39d617b4bfe333d3946b2bb..86e755c22963c2b27605d8e2a06fcd998d8b4aa1 100644 GIT binary patch delta 870 zcmah{&rj1}7;eXY*$r{XN?8^)Q5s0pB1rsk@@-c-l(DXnLU)kNGG-9jWaHOjvWq8T zG~x5$$*YOsVB~5d2Tcr#iGP9KJQ=v~2l&42M571av~RzD@AEw0^SOOdJ2^W*YvC;JqckGwb` zzwSy)>~qIiw%0Mp4nwJP;?r(XkfY-F1lO0A;$T_%>Fd5LX*aA;472yq8?7o}=uiPz z#|Mm8Q~=e10>CYLH?aWPWdfBc;nuG_XUtLMs>($m?;dj8ZN0C6Kz!;mdrM@aG%cqcf^Fx zTD7~=_B@IDBX*8Pj8)D5WSOb-aC=H%GoVZJ$V)lg`vi%MVw|jSFs#y^Zzq+w4m$ T$-c$H@ifPAAEO*+f_3R9jSKaW delta 736 zcmah`y-(Xv5O+unCLw@8fYW3^s$_u>GC_Q#NJ!jE+{C2`jR^#23kDfT!dHQ`EMchB zy>E@G4jnq6sug>t{t2Disaw^pOQ+sFn~I?lH{AX1``+(;dGKrO)0?gvC3qT$Ue8Nv z_j6}&w^1D#;;%!mxF*|y_RW=C?c!@nO#T+u-iCQY+35Mur+tkiRQ@?Q%s&QW{I`;l zulu!5l*dB_e$^X)_-=4CtSO^^2BcU*?E}kND=&Go)t|CX@{)IYS6Vdy$QV!r*aHK= za2W<{@)r=mQx_GRBZwhiv!;Hb>*z4{LIkM@sL_`!*-CcvAYU@6*92%%9<3;CP@WNN ztsYG2dJ2(|!}WfYF1eKr3f3y+{A__j`9(ZEwN6kh6qOn0o`W7!y*TL7gtSXKwaLyE+Ys6~Y&xB;%- X*|ft$Q8m!! { + thrift_write_field_bool(&mut w, &mut inner, 1, *utc); + thrift_write_field_struct_begin(&mut w, &mut inner, 2); + let mut unit_last = 0i16; + thrift_write_field_struct_begin(&mut w, &mut unit_last, time_unit_union_id(logical)); + thrift_write_field_stop(&mut w); + thrift_write_field_stop(&mut w); + } + // Time of day has no zone, so isAdjustedToUTC is false. + #[cfg(feature = "datetime")] + ParquetLogicalType::TimeMillis | ParquetLogicalType::TimeMicros | ParquetLogicalType::TimeNanos => { thrift_write_field_bool(&mut w, &mut inner, 1, false); @@ -502,9 +511,9 @@ fn logical_type_union_id(logical: &ParquetLogicalType) -> Option { | ParquetLogicalType::TimeMicros | ParquetLogicalType::TimeNanos => 7, #[cfg(feature = "datetime")] - ParquetLogicalType::TimestampMillis - | ParquetLogicalType::TimestampMicros - | ParquetLogicalType::TimestampNanos => 8, + ParquetLogicalType::TimestampMillis { .. } + | ParquetLogicalType::TimestampMicros { .. } + | ParquetLogicalType::TimestampNanos { .. } => 8, ParquetLogicalType::IntType { .. } => 10, }) } @@ -513,8 +522,8 @@ fn logical_type_union_id(logical: &ParquetLogicalType) -> Option { #[cfg(feature = "datetime")] fn time_unit_union_id(logical: &ParquetLogicalType) -> i16 { match logical { - ParquetLogicalType::TimestampMillis | ParquetLogicalType::TimeMillis => 1, - ParquetLogicalType::TimestampMicros | ParquetLogicalType::TimeMicros => 2, + ParquetLogicalType::TimestampMillis { .. } | ParquetLogicalType::TimeMillis => 1, + ParquetLogicalType::TimestampMicros { .. } | ParquetLogicalType::TimeMicros => 2, _ => 3, } } diff --git a/rust/src/models/readers/parquet.rs b/rust/src/models/readers/parquet.rs index a33f6f1..c461e20 100644 --- a/rust/src/models/readers/parquet.rs +++ b/rust/src/models/readers/parquet.rs @@ -1034,12 +1034,14 @@ fn parse_logical_type(r: &mut R) -> Result, 7 | 8 => { let mut inner = 0i16; let mut unit = 0i16; + let mut utc = false; loop { let (t, f) = thrift_read_field_begin(r, &mut inner)?; if t == 0 { break; } match f { + 1 if t == TC_BOOL_TRUE || t == TC_BOOL_FALSE => utc = t == TC_BOOL_TRUE, 2 if t == TC_STRUCT => { let mut unit_last = 0i16; loop { @@ -1058,9 +1060,9 @@ fn parse_logical_type(r: &mut R) -> Result, (7, 1) => Some(ParquetLogicalType::TimeMillis), (7, 2) => Some(ParquetLogicalType::TimeMicros), (7, 3) => Some(ParquetLogicalType::TimeNanos), - (8, 1) => Some(ParquetLogicalType::TimestampMillis), - (8, 2) => Some(ParquetLogicalType::TimestampMicros), - (8, 3) => Some(ParquetLogicalType::TimestampNanos), + (8, 1) => Some(ParquetLogicalType::TimestampMillis { utc }), + (8, 2) => Some(ParquetLogicalType::TimestampMicros { utc }), + (8, 3) => Some(ParquetLogicalType::TimestampNanos { utc }), _ => None, } } diff --git a/rust/src/models/types/parquet.rs b/rust/src/models/types/parquet.rs index 0f8be28..8a932a9 100644 --- a/rust/src/models/types/parquet.rs +++ b/rust/src/models/types/parquet.rs @@ -73,15 +73,19 @@ pub(crate) enum ParquetLogicalType { /// DATE - days since the Unix epoch, stored as INT32. #[cfg(feature = "datetime")] Date32, - /// 64-bit timestamp - milliseconds since epoch + /// 64-bit timestamp - milliseconds since epoch. `utc` is Parquet's + /// `isAdjustedToUTC`: true for instants in UTC, false for local + /// wall-clock time with no zone. #[cfg(feature = "datetime")] - TimestampMillis, - /// 64-bit timestamp - microseconds since epoch + TimestampMillis { utc: bool }, + /// 64-bit timestamp - microseconds since epoch, with the same `utc` + /// meaning as `TimestampMillis`. #[cfg(feature = "datetime")] - TimestampMicros, - /// 64-bit timestamp - nanoseconds since epoch + TimestampMicros { utc: bool }, + /// 64-bit timestamp - nanoseconds since epoch, with the same `utc` + /// meaning as `TimestampMillis`. Exists only as a `LogicalType`. #[cfg(feature = "datetime")] - TimestampNanos, + TimestampNanos { utc: bool }, /// 32-bit time - milliseconds since midnight #[cfg(feature = "datetime")] TimeMillis, @@ -136,10 +140,11 @@ impl ParquetLogicalType { Some(7) => Some(ParquetLogicalType::TimeMillis), #[cfg(feature = "datetime")] Some(8) => Some(ParquetLogicalType::TimeMicros), + // The legacy TIMESTAMP converted types are defined as UTC instants. #[cfg(feature = "datetime")] - Some(9) => Some(ParquetLogicalType::TimestampMillis), + Some(9) => Some(ParquetLogicalType::TimestampMillis { utc: true }), #[cfg(feature = "datetime")] - Some(10) => Some(ParquetLogicalType::TimestampMicros), + Some(10) => Some(ParquetLogicalType::TimestampMicros { utc: true }), Some(11) => Some(ParquetLogicalType::IntType { bit_width: 8, is_signed: false, @@ -325,21 +330,24 @@ pub(crate) fn arrow_type_to_parquet( // carried into days on write. See `temporal_unit_scale`. #[cfg(feature = "datetime")] ArrowType::Date64 => Ok((ParquetPhysicalType::Int32, ParquetLogicalType::Date32)), + // A timestamp with a timezone is an instant, so it is stored as + // adjusted to UTC. A timestamp without one is local wall-clock time. + // Parquet keeps only that flag, not the zone name. #[cfg(feature = "datetime")] - ArrowType::Timestamp(unit, _) => match unit { + ArrowType::Timestamp(unit, tz) => match unit { // Parquet has no seconds unit, so seconds are scaled to // milliseconds on write. TimeUnit::Seconds | TimeUnit::Milliseconds => Ok(( ParquetPhysicalType::Int64, - ParquetLogicalType::TimestampMillis, + ParquetLogicalType::TimestampMillis { utc: tz.is_some() }, )), TimeUnit::Microseconds => Ok(( ParquetPhysicalType::Int64, - ParquetLogicalType::TimestampMicros, + ParquetLogicalType::TimestampMicros { utc: tz.is_some() }, )), TimeUnit::Nanoseconds => Ok(( ParquetPhysicalType::Int64, - ParquetLogicalType::TimestampNanos, + ParquetLogicalType::TimestampNanos { utc: tz.is_some() }, )), // A timestamp counted in days is a date. TimeUnit::Days => Ok((ParquetPhysicalType::Int32, ParquetLogicalType::Date32)), @@ -472,18 +480,20 @@ pub(crate) fn parquet_to_arrow_type( // Dates, times, timestamps #[cfg(feature = "datetime")] (ParquetPhysicalType::Int32, Some(ParquetLogicalType::Date32)) => Ok(ArrowType::Date32), + // A UTC-adjusted timestamp reads back with the "UTC" zone. Parquet + // records no zone name, so that is the only zone a file can carry. #[cfg(feature = "datetime")] - (ParquetPhysicalType::Int64, Some(ParquetLogicalType::TimestampMillis)) => { - Ok(ArrowType::Timestamp(TimeUnit::Milliseconds, None)) - } + (ParquetPhysicalType::Int64, Some(ParquetLogicalType::TimestampMillis { utc })) => Ok( + ArrowType::Timestamp(TimeUnit::Milliseconds, utc.then(|| "UTC".to_string())), + ), #[cfg(feature = "datetime")] - (ParquetPhysicalType::Int64, Some(ParquetLogicalType::TimestampMicros)) => { - Ok(ArrowType::Timestamp(TimeUnit::Microseconds, None)) - } + (ParquetPhysicalType::Int64, Some(ParquetLogicalType::TimestampMicros { utc })) => Ok( + ArrowType::Timestamp(TimeUnit::Microseconds, utc.then(|| "UTC".to_string())), + ), #[cfg(feature = "datetime")] - (ParquetPhysicalType::Int64, Some(ParquetLogicalType::TimestampNanos)) => { - Ok(ArrowType::Timestamp(TimeUnit::Nanoseconds, None)) - } + (ParquetPhysicalType::Int64, Some(ParquetLogicalType::TimestampNanos { utc })) => Ok( + ArrowType::Timestamp(TimeUnit::Nanoseconds, utc.then(|| "UTC".to_string())), + ), #[cfg(feature = "datetime")] (ParquetPhysicalType::Int32, Some(ParquetLogicalType::TimeMillis)) => { Ok(ArrowType::Time32(TimeUnit::Milliseconds)) diff --git a/rust/src/models/writers/parquet.rs b/rust/src/models/writers/parquet.rs index 191fa70..bbf5cd6 100644 --- a/rust/src/models/writers/parquet.rs +++ b/rust/src/models/writers/parquet.rs @@ -659,10 +659,12 @@ fn logical_to_converted(log: &ParquetLogicalType) -> Option { ParquetLogicalType::Utf8 => 0, #[cfg(feature = "datetime")] ParquetLogicalType::Date32 => 6, + // The legacy TIMESTAMP converted types mean UTC instants, so a local + // timestamp carries only its LogicalType. #[cfg(feature = "datetime")] - ParquetLogicalType::TimestampMillis => 9, + ParquetLogicalType::TimestampMillis { utc: true } => 9, #[cfg(feature = "datetime")] - ParquetLogicalType::TimestampMicros => 10, + ParquetLogicalType::TimestampMicros { utc: true } => 10, #[cfg(feature = "datetime")] ParquetLogicalType::TimeMillis => 7, #[cfg(feature = "datetime")] diff --git a/rust/tests/parquet_nullable_roundtrip.rs b/rust/tests/parquet_nullable_roundtrip.rs index 1dd935a..f1adff1 100644 --- a/rust/tests/parquet_nullable_roundtrip.rs +++ b/rust/tests/parquet_nullable_roundtrip.rs @@ -266,6 +266,12 @@ mod parquet_nullable_roundtrip_tests { TimeUnit::Nanoseconds, timestamp_value, )); + cols.push(temporal64( + "ts_us_utc", + ArrowType::Timestamp(TimeUnit::Microseconds, Some("UTC".to_string())), + TimeUnit::Microseconds, + timestamp_value, + )); cols.push(temporal32( "time32_ms", ArrowType::Time32(TimeUnit::Milliseconds), @@ -475,6 +481,15 @@ mod parquet_nullable_roundtrip_tests { TimeUnit::Nanoseconds, timestamp_value, ); + // A zoned timestamp is stored as adjusted to UTC and reads back + // with the UTC zone. + assert_temporal64( + out, + "ts_us_utc", + ArrowType::Timestamp(TimeUnit::Microseconds, Some("UTC".to_string())), + TimeUnit::Microseconds, + timestamp_value, + ); assert_temporal32( out, "time32_ms", diff --git a/rust/tests/pyarrow_parquet.rs b/rust/tests/pyarrow_parquet.rs index 84dde80..4401acd 100644 --- a/rust/tests/pyarrow_parquet.rs +++ b/rust/tests/pyarrow_parquet.rs @@ -267,6 +267,15 @@ mod pyarrow_parquet_tests { temporal64(&table, "ts_ns", TimeUnit::Nanoseconds), [Some(0), Some(1_700_000_000_000_000_000), None, Some(-1), Some(1)] ); + // pyarrow writes a zoned timestamp with isAdjustedToUTC set. + assert_eq!( + column_dtype(&table, "ts_us_utc"), + ArrowType::Timestamp(TimeUnit::Microseconds, Some("UTC".to_string())) + ); + assert_eq!( + temporal64(&table, "ts_us_utc", TimeUnit::Microseconds), + [Some(0), Some(1_700_000_000_000_000), None, Some(-1), Some(1)] + ); assert_eq!(column_dtype(&table, "dec32"), ArrowType::Decimal32(7, 2)); match column(&table, "dec32") {