From 19002abc3522af60b6eaa99c0d5d4999b70cd681 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 14 Apr 2026 20:02:22 +0300 Subject: [PATCH 1/9] fix: correct DictEncoder::estimated_memory_size The returned value should estimate the actual memory usage, but instead it used the evaluation of the encoded size of the dictionary data, and bypassed the hash table memory usage added by the Interner. The implementation of Storage::estimated_memory_size for the unique key storage was not correct as well, but it was unused. Correct both problems. --- parquet/src/encodings/encoding/dict_encoder.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index 79a1f247670c..15c61a3e7971 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -64,7 +64,7 @@ impl Storage for KeyStorage { } fn estimated_memory_size(&self) -> usize { - self.size_in_bytes + self.uniques.capacity() * std::mem::size_of::() + self.uniques.capacity() * std::mem::size_of::() } } @@ -183,6 +183,6 @@ impl Encoder for DictEncoder { /// /// For this encoder, the indices are unencoded bytes (refer to [`Self::write_indices`]). fn estimated_memory_size(&self) -> usize { - self.interner.storage().size_in_bytes + self.indices.len() * std::mem::size_of::() + self.interner.estimated_memory_size() + self.indices.len() * std::mem::size_of::() } } From 70b09123c6658b40d0a206c9f524ea1810c11f03 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 18:54:53 +0300 Subject: [PATCH 2/9] fix(parquet): Interner::estimated_memory_size Should multiply hash table capacity by the key size, not add them. --- parquet/src/util/interner.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parquet/src/util/interner.rs b/parquet/src/util/interner.rs index 34c7d1390f7a..bf4374fc43b5 100644 --- a/parquet/src/util/interner.rs +++ b/parquet/src/util/interner.rs @@ -79,7 +79,7 @@ impl Interner { pub fn estimated_memory_size(&self) -> usize { self.storage.estimated_memory_size() + // estimate size of dedup hashmap as just th size of the keys - self.dedup.capacity() + std::mem::size_of::() + self.dedup.capacity() * std::mem::size_of::() } /// Returns the storage for this interner From ffa40d98a1cbc3a7c3906da8a27824d16c4aaf4b Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 19:22:38 +0300 Subject: [PATCH 3/9] fix(parquet): use HashTable::allocation_size --- parquet/src/util/interner.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/parquet/src/util/interner.rs b/parquet/src/util/interner.rs index bf4374fc43b5..deae3720d5cf 100644 --- a/parquet/src/util/interner.rs +++ b/parquet/src/util/interner.rs @@ -77,9 +77,7 @@ impl Interner { /// Return estimate of the memory used, in bytes #[allow(dead_code)] // not used in parquet_derive, so is dead there pub fn estimated_memory_size(&self) -> usize { - self.storage.estimated_memory_size() + - // estimate size of dedup hashmap as just th size of the keys - self.dedup.capacity() * std::mem::size_of::() + self.storage.estimated_memory_size() + self.dedup.allocation_size() } /// Returns the storage for this interner From 07aba13a058cd886c05b3c546c4a607a949c57de Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 20:20:20 +0300 Subject: [PATCH 4/9] fix: account for allocated bytes in KeyStorage --- parquet/src/encodings/encoding/dict_encoder.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index 15c61a3e7971..7d38c9739ecc 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -64,7 +64,12 @@ impl Storage for KeyStorage { } fn estimated_memory_size(&self) -> usize { - self.uniques.capacity() * std::mem::size_of::() + let uniques_heap_bytes = match T::get_physical_type() { + Type::FIXED_LEN_BYTE_ARRAY => self.type_length * self.uniques.len(), + _ => ::variable_length_bytes(&self.uniques) + .unwrap_or(0) as usize, + }; + self.uniques.capacity() * std::mem::size_of::() + uniques_heap_bytes } } From bf6a03a3111e0822cd6c746734f6a617d5663ecb Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 20:23:08 +0300 Subject: [PATCH 5/9] fix: use vector capacity, not length To estimate allocation size, the vector capacity is the right thing to use, not the length. --- parquet/src/encodings/encoding/dict_encoder.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index 7d38c9739ecc..af64b99527df 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -188,6 +188,7 @@ impl Encoder for DictEncoder { /// /// For this encoder, the indices are unencoded bytes (refer to [`Self::write_indices`]). fn estimated_memory_size(&self) -> usize { - self.interner.estimated_memory_size() + self.indices.len() * std::mem::size_of::() + self.interner.estimated_memory_size() + + self.indices.capacity() * std::mem::size_of::() } } From 52795a8b003be832398be32c019c4feaf00394ca Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 20:24:08 +0300 Subject: [PATCH 6/9] test: DictEncoder::estimated_memory_size tests As we cannot test the exact memory allocation behavior or even give the upper bound without delving into implementation details of Vec and hashbrown, the only reliable test is to confirm that the lower bounds are respected, that is, increases in memory use from the empty state (where all vectors are empty and therefore have added no allocations) add at least the expected amount of memory. --- .../src/encodings/encoding/dict_encoder.rs | 205 ++++++++++++++++++ 1 file changed, 205 insertions(+) diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index af64b99527df..a96afb74508d 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -192,3 +192,208 @@ impl Encoder for DictEncoder { + self.indices.capacity() * std::mem::size_of::() } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::data_type::{ByteArray, ByteArrayType, FixedLenByteArray, FixedLenByteArrayType, Int32Type}; + use crate::encodings::encoding::Encoder; + use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType}; + + fn make_col_desc() -> ColumnDescPtr { + make_col_desc_with_length::(-1) + } + + fn make_col_desc_with_length(type_length: i32) -> ColumnDescPtr { + let ty = SchemaType::primitive_type_builder("col", T::get_physical_type()) + .with_length(type_length) + .build() + .unwrap(); + Arc::new(ColumnDescriptor::new( + Arc::new(ty), + 0, + 0, + ColumnPath::new(vec![]), + )) + } + + #[test] + fn test_estimated_memory_size_primitive_with_duplicates() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + let empty_size = encoder.estimated_memory_size(); + + // 3 distinct values, repeated to produce 9 indices total. + encoder.put(&[1, 2, 3, 1, 2, 3, 1, 2, 3]).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 3 unique dictionary entries. + let dict_entry_size = 3 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 9 buffered indices. + let indices_size = 9 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_primitive_all_distinct() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + let empty_size = encoder.estimated_memory_size(); + + let values: Vec = (0..100).collect(); + encoder.put(&values).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 100 unique dictionary entries. + let dict_entry_size = 100 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 100 buffered indices. + let indices_size = 100 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_byte_array_with_duplicates() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + let empty_size = encoder.estimated_memory_size(); + + // 3 distinct byte strings ("foo", "bar", "baz" — 3 bytes each), repeated to produce + // 9 indices total. + let vals: Vec = [ + "foo", "bar", "baz", "foo", "bar", "baz", "foo", "bar", "baz", + ] + .iter() + .map(|s| ByteArray::from(*s)) + .collect(); + encoder.put(&vals).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 3 unique dictionary entries, including their heap-allocated bytes. + let dict_entry_size = 3 * std::mem::size_of::() + 3 * 3; // 3 values × 3 bytes each + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 9 buffered indices. + let indices_size = 9 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_byte_array_all_distinct() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + let empty_size = encoder.estimated_memory_size(); + + // 100 distinct values: "0".."9" (1 byte each) and "10".."99" (2 bytes each). + let values: Vec = (0..100_u32) + .map(|i| ByteArray::from(i.to_string().into_bytes())) + .collect(); + let bytes_total: usize = values.iter().map(|v| v.len()).sum(); // 10×1 + 90×2 = 190 + encoder.put(&values).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 100 unique dictionary entries, including their heap-allocated bytes. + let dict_entry_size = 100 * std::mem::size_of::() + bytes_total; + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 100 buffered indices. + let indices_size = 100 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_fixed_len_byte_array_with_duplicates() { + const TYPE_LEN: usize = 3; + let mut encoder = DictEncoder::::new( + make_col_desc_with_length::(TYPE_LEN as i32), + ); + let empty_size = encoder.estimated_memory_size(); + + // 3 distinct 3-byte values, repeated to produce 9 indices total. + let vals: Vec = [ + b"foo", b"bar", b"baz", b"foo", b"bar", b"baz", b"foo", b"bar", b"baz", + ] + .iter() + .map(|b| FixedLenByteArray::from(b.to_vec())) + .collect(); + encoder.put(&vals).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 3 unique dictionary entries: struct overhead plus the + // fixed-length bytes allocated per entry. + let dict_entry_size = 3 * std::mem::size_of::() + 3 * TYPE_LEN; + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 9 buffered indices. + let indices_size = 9 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } + + #[test] + fn test_estimated_memory_size_fixed_len_byte_array_all_distinct() { + const TYPE_LEN: usize = 3; + let mut encoder = DictEncoder::::new( + make_col_desc_with_length::(TYPE_LEN as i32), + ); + let empty_size = encoder.estimated_memory_size(); + + // 100 distinct 3-byte values: zero-padded big-endian u8 indices. + let values: Vec = (0..100_u8) + .map(|i| FixedLenByteArray::from(vec![0u8, 0u8, i])) + .collect(); + encoder.put(&values).unwrap(); + + let size = encoder.estimated_memory_size(); + + // Must account for the 100 unique dictionary entries: struct overhead plus the + // fixed-length bytes allocated per entry. + let dict_entry_size = 100 * std::mem::size_of::() + 100 * TYPE_LEN; + assert!( + size >= empty_size + dict_entry_size, + "memory size {size} should grow by at least the dict storage ({dict_entry_size} bytes)" + ); + + // Must also account for the 100 buffered indices. + let indices_size = 100 * std::mem::size_of::(); + assert!( + size >= empty_size + dict_entry_size + indices_size, + "memory size {size} should include indices ({indices_size} bytes)" + ); + } +} From 986b08490693ec30737ab92d28ac2c57ace0e7c9 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 20:36:15 +0300 Subject: [PATCH 7/9] chore: rustfmt 1.95.0 --- .../src/encodings/encoding/dict_encoder.rs | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index a96afb74508d..37cfdb9ba155 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -198,7 +198,9 @@ mod tests { use std::sync::Arc; use super::*; - use crate::data_type::{ByteArray, ByteArrayType, FixedLenByteArray, FixedLenByteArrayType, Int32Type}; + use crate::data_type::{ + ByteArray, ByteArrayType, FixedLenByteArray, FixedLenByteArrayType, Int32Type, + }; use crate::encodings::encoding::Encoder; use crate::schema::types::{ColumnDescriptor, ColumnPath, Type as SchemaType}; @@ -333,18 +335,18 @@ mod tests { #[test] fn test_estimated_memory_size_fixed_len_byte_array_with_duplicates() { const TYPE_LEN: usize = 3; - let mut encoder = DictEncoder::::new( - make_col_desc_with_length::(TYPE_LEN as i32), - ); + let mut encoder = DictEncoder::::new(make_col_desc_with_length::< + FixedLenByteArrayType, + >(TYPE_LEN as i32)); let empty_size = encoder.estimated_memory_size(); // 3 distinct 3-byte values, repeated to produce 9 indices total. - let vals: Vec = [ + let vals = [ b"foo", b"bar", b"baz", b"foo", b"bar", b"baz", b"foo", b"bar", b"baz", ] .iter() .map(|b| FixedLenByteArray::from(b.to_vec())) - .collect(); + .collect::>(); encoder.put(&vals).unwrap(); let size = encoder.estimated_memory_size(); @@ -368,15 +370,15 @@ mod tests { #[test] fn test_estimated_memory_size_fixed_len_byte_array_all_distinct() { const TYPE_LEN: usize = 3; - let mut encoder = DictEncoder::::new( - make_col_desc_with_length::(TYPE_LEN as i32), - ); + let mut encoder = DictEncoder::::new(make_col_desc_with_length::< + FixedLenByteArrayType, + >(TYPE_LEN as i32)); let empty_size = encoder.estimated_memory_size(); // 100 distinct 3-byte values: zero-padded big-endian u8 indices. - let values: Vec = (0..100_u8) + let values = (0..100_u8) .map(|i| FixedLenByteArray::from(vec![0u8, 0u8, i])) - .collect(); + .collect::>(); encoder.put(&values).unwrap(); let size = encoder.estimated_memory_size(); From e454a9c9a6482ec995e969d61b37125d46a597cb Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 24 Apr 2026 05:22:40 +0300 Subject: [PATCH 8/9] More DictEncoder::estimated_memory_size tests Add tests demonstrating failures of the previous implementation, also one asserting that the uniques vector's capacity is included in the estimate. --- .../src/encodings/encoding/dict_encoder.rs | 68 +++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index 37cfdb9ba155..2248066a50cc 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -398,4 +398,72 @@ mod tests { "memory size {size} should include indices ({indices_size} bytes)" ); } + + #[test] + fn test_estimated_memory_size_includes_interner_dedup_table() { + // The dedup `HashTable` in `Interner` is preallocated with + // `DEFAULT_DEDUP_CAPACITY` slots at construction, independent of any + // values pushed. + let encoder = DictEncoder::::new(make_col_desc::()); + + let size = encoder.estimated_memory_size(); + + assert!( + size > 0, + "memory size should include the preallocated dedup hash table" + ); + } + + #[test] + fn test_estimated_memory_size_accounts_for_indices_capacity() { + // Exercises the `indices.capacity()` (not `.len()`) accounting. + // After a flush, `indices` is cleared but its capacity is retained; pushing a + // smaller batch afterwards leaves capacity strictly greater than length. + let mut encoder = DictEncoder::::new(make_col_desc::()); + + let big: Vec = vec![0; 64]; + encoder.put(&big).unwrap(); + let _ = encoder.flush_buffer().unwrap(); + + let flushed_size = encoder.estimated_memory_size(); + + // Push a single value — indices.len() == 1 but indices.capacity() >= 64. + // No change on the key storage since the value is already interned. + encoder.put(&[0]).unwrap(); + + let size = encoder.estimated_memory_size(); + + assert_eq!( + size, flushed_size, + "memory size should include retained indices capacity", + ); + } + + #[test] + fn test_estimated_memory_size_accounts_for_uniques_capacity() { + let mut encoder = DictEncoder::::new(make_col_desc::()); + + let values: Vec = (0..64).collect(); + encoder.put(&values).unwrap(); + // Flush indices so they don't mask the uniques accounting in the lower bound. + let _ = encoder.flush_buffer().unwrap(); + + let size1 = encoder.estimated_memory_size(); + + // Push more values to trigger uniques capacity growth. + // The pre-allocated dedup hash table is unlikely to be resized. + let values: Vec = (64..128).collect(); + encoder.put(&values).unwrap(); + // Flush indices so they don't mask the uniques accounting in the lower bound. + let _ = encoder.flush_buffer().unwrap(); + + let size2 = encoder.estimated_memory_size(); + + let min_uniques_bytes = 64 * std::mem::size_of::(); + assert!( + size2 >= size1 + min_uniques_bytes, + "memory size {size2} should grow from {size1} by allocated uniques capacity \ + (at least {min_uniques_bytes} bytes)" + ); + } } From 3dd032af11bfb917467da37faa600b98df720f1b Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 27 Apr 2026 21:24:41 +0300 Subject: [PATCH 9/9] fix: memory size for DictEncoder indices The elements of the indices vector are u64, so the memory size should be calculated accordingly. --- parquet/src/encodings/encoding/dict_encoder.rs | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index 2248066a50cc..89666bbe7313 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -188,8 +188,7 @@ impl Encoder for DictEncoder { /// /// For this encoder, the indices are unencoded bytes (refer to [`Self::write_indices`]). fn estimated_memory_size(&self) -> usize { - self.interner.estimated_memory_size() - + self.indices.capacity() * std::mem::size_of::() + self.interner.estimated_memory_size() + self.indices.capacity() * std::mem::size_of::() } } @@ -239,7 +238,7 @@ mod tests { ); // Must also account for the 9 buffered indices. - let indices_size = 9 * std::mem::size_of::(); + let indices_size = 9 * std::mem::size_of::(); assert!( size >= empty_size + dict_entry_size + indices_size, "memory size {size} should include indices ({indices_size} bytes)" @@ -264,7 +263,7 @@ mod tests { ); // Must also account for the 100 buffered indices. - let indices_size = 100 * std::mem::size_of::(); + let indices_size = 100 * std::mem::size_of::(); assert!( size >= empty_size + dict_entry_size + indices_size, "memory size {size} should include indices ({indices_size} bytes)" @@ -296,7 +295,7 @@ mod tests { ); // Must also account for the 9 buffered indices. - let indices_size = 9 * std::mem::size_of::(); + let indices_size = 9 * std::mem::size_of::(); assert!( size >= empty_size + dict_entry_size + indices_size, "memory size {size} should include indices ({indices_size} bytes)" @@ -325,7 +324,7 @@ mod tests { ); // Must also account for the 100 buffered indices. - let indices_size = 100 * std::mem::size_of::(); + let indices_size = 100 * std::mem::size_of::(); assert!( size >= empty_size + dict_entry_size + indices_size, "memory size {size} should include indices ({indices_size} bytes)" @@ -360,7 +359,7 @@ mod tests { ); // Must also account for the 9 buffered indices. - let indices_size = 9 * std::mem::size_of::(); + let indices_size = 9 * std::mem::size_of::(); assert!( size >= empty_size + dict_entry_size + indices_size, "memory size {size} should include indices ({indices_size} bytes)" @@ -392,7 +391,7 @@ mod tests { ); // Must also account for the 100 buffered indices. - let indices_size = 100 * std::mem::size_of::(); + let indices_size = 100 * std::mem::size_of::(); assert!( size >= empty_size + dict_entry_size + indices_size, "memory size {size} should include indices ({indices_size} bytes)"