This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-mosaic.git
The following commit(s) were added to refs/heads/main by this push:
new cf6a330 perf(core): reduce dictionary and nullable write overhead
(#91)
cf6a330 is described below
commit cf6a330aee7f76d2f0bd1c4694e2a24520cdf084
Author: jianguotian <[email protected]>
AuthorDate: Sun Sep 20 20:03:41 2026 +0800
perf(core): reduce dictionary and nullable write overhead (#91)
---
core/src/bucket_writer.rs | 497 ++++++++++++++++++++++++++++++++++++----------
1 file changed, 396 insertions(+), 101 deletions(-)
diff --git a/core/src/bucket_writer.rs b/core/src/bucket_writer.rs
index 487a1a3..9828c90 100644
--- a/core/src/bucket_writer.rs
+++ b/core/src/bucket_writer.rs
@@ -54,6 +54,13 @@ pub struct ChildColumnMeta {
pub num_elements: usize,
}
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum DictTracking {
+ Pending,
+ Active,
+ Disabled,
+}
+
pub struct BucketWriter {
num_primary: usize,
total_columns: usize,
@@ -68,6 +75,7 @@ pub struct BucketWriter {
long_dict_maps: Vec<Option<HashMap<u64, usize>>>,
byte_dict_maps: Vec<Option<HashMap<Vec<u8>, usize>>>,
+ dict_tracking: Vec<DictTracking>,
dict_total_bytes: Vec<usize>,
max_dict_total_bytes: usize,
max_dict_entries: usize,
@@ -87,18 +95,6 @@ impl BucketWriter {
let total_columns = physical_types.len();
let fixed_widths: Vec<i32> =
physical_types.iter().map(types::fixed_width).collect();
- let mut long_dict_maps = Vec::with_capacity(total_columns);
- let mut byte_dict_maps = Vec::with_capacity(total_columns);
- for fw in &fixed_widths {
- if uses_long_dict(*fw) {
- long_dict_maps.push(Some(HashMap::new()));
- byte_dict_maps.push(None);
- } else {
- long_dict_maps.push(None);
- byte_dict_maps.push(Some(HashMap::new()));
- }
- }
-
BucketWriter {
num_primary,
total_columns,
@@ -108,8 +104,9 @@ impl BucketWriter {
non_null_counts: vec![0; total_columns],
const_tracking: vec![true; total_columns],
first_value_len: vec![0; total_columns],
- long_dict_maps,
- byte_dict_maps,
+ long_dict_maps: (0..total_columns).map(|_| None).collect(),
+ byte_dict_maps: (0..total_columns).map(|_| None).collect(),
+ dict_tracking: vec![DictTracking::Pending; total_columns],
dict_total_bytes: vec![0; total_columns],
max_dict_total_bytes,
max_dict_entries,
@@ -318,8 +315,17 @@ impl BucketWriter {
}
let typed = downcast_array(array, dt)?;
+ let null_count = array.null_count();
+
+ if null_count == num_new_rows {
+ let nulls = array
+ .nulls()
+ .expect("an all-null array must have a physical null buffer");
+ append_null_bitmap(&mut self.null_bitmaps[col], start_row, nulls);
+ return Ok(0);
+ }
- if array.null_count() == 0 {
+ if null_count == 0 {
if let Some(col_size) =
self.append_no_null_batch(col, &typed, start_row,
num_new_rows)?
{
@@ -327,6 +333,15 @@ impl BucketWriter {
}
}
+ if let TypedArrayRef::TimestampMillis(array) = &typed {
+ return Ok(self.append_nullable_timestamp_millis_batch(
+ col,
+ array,
+ start_row,
+ num_new_rows,
+ ));
+ }
+
let mut col_size = 0usize;
for row in 0..num_new_rows {
let abs_row = start_row + row;
@@ -339,40 +354,7 @@ impl BucketWriter {
let written = self.value_buffers[col].len() - before;
col_size += written;
- if self.const_tracking[col] {
- if self.non_null_counts[col] == 1 {
- self.first_value_len[col] = written;
- } else if written != self.first_value_len[col]
- || !equals_first_value(&self.value_buffers[col],
before, written)
- {
- self.const_tracking[col] = false;
- }
- }
-
- if let Some(ref mut dict) = self.long_dict_maps[col] {
- let key = values::extract_fixed_key(
- &self.value_buffers[col],
- before,
- self.fixed_widths[col],
- );
- let len = dict.len();
- dict.entry(key).or_insert(len);
- if dict.len() > self.max_dict_entries {
- self.long_dict_maps[col] = None;
- }
- } else if let Some(ref mut dict) = self.byte_dict_maps[col] {
- let slice = &self.value_buffers[col][before..before +
written];
- if !dict.contains_key(slice) {
- let len = dict.len();
- dict.insert(slice.to_vec(), len);
- self.dict_total_bytes[col] += written;
- }
- if dict.len() > self.max_dict_entries
- || self.dict_total_bytes[col] >
self.max_dict_total_bytes
- {
- self.byte_dict_maps[col] = None;
- }
- }
+ self.track_encoding_value(col, before, written);
}
}
Ok(col_size)
@@ -492,52 +474,14 @@ impl BucketWriter {
let col_size = buf.len() - before_all;
let fw = self.fixed_widths[col] as usize;
- let prev_non_null = self.non_null_counts[col];
- self.non_null_counts[col] += num_rows;
-
- if self.const_tracking[col] {
- if prev_non_null == 0 {
- self.first_value_len[col] = fw;
- for i in 1..num_rows {
- if !equals_first_value(buf, before_all + i * fw, fw) {
- self.const_tracking[col] = false;
- break;
- }
- }
- } else {
- for i in 0..num_rows {
- if !equals_first_value(buf, before_all + i * fw, fw) {
- self.const_tracking[col] = false;
- break;
- }
- }
- }
- }
-
- if let Some(ref mut dict) = self.long_dict_maps[col] {
- for i in 0..num_rows {
- let key =
- values::extract_fixed_key(buf, before_all + i * fw,
self.fixed_widths[col]);
- let len = dict.len();
- dict.entry(key).or_insert(len);
- if dict.len() > self.max_dict_entries {
- self.long_dict_maps[col] = None;
- break;
- }
- }
- } else if let Some(ref mut dict) = self.byte_dict_maps[col] {
+ if !self.const_tracking[col] && self.dict_tracking[col] ==
DictTracking::Disabled {
+ self.non_null_counts[col] += num_rows;
+ } else {
for i in 0..num_rows {
- let start = before_all + i * fw;
- let slice = &buf[start..start + fw];
- if !dict.contains_key(slice) {
- let len = dict.len();
- dict.insert(slice.to_vec(), len);
- self.dict_total_bytes[col] += fw;
- }
- if dict.len() > self.max_dict_entries
- || self.dict_total_bytes[col] > self.max_dict_total_bytes
- {
- self.byte_dict_maps[col] = None;
+ self.non_null_counts[col] += 1;
+ self.track_encoding_value(col, before_all + i * fw, fw);
+ if !self.const_tracking[col] && self.dict_tracking[col] ==
DictTracking::Disabled {
+ self.non_null_counts[col] += num_rows - i - 1;
break;
}
}
@@ -549,6 +493,116 @@ impl BucketWriter {
Ok(Some(col_size))
}
+ fn append_nullable_timestamp_millis_batch(
+ &mut self,
+ col: usize,
+ array: &TimestampMillisecondArray,
+ start_row: usize,
+ num_rows: usize,
+ ) -> usize {
+ let nulls = array
+ .nulls()
+ .expect("nullable timestamp batch must have a physical null
buffer");
+ debug_assert!(nulls.null_count() > 0);
+ debug_assert!(nulls.null_count() < num_rows);
+
+ append_null_bitmap(&mut self.null_bitmaps[col], start_row, nulls);
+
+ let valid_count = num_rows - nulls.null_count();
+ self.value_buffers[col].reserve(valid_count * 8);
+ let values = array.values();
+ let mut track_encoding =
+ self.const_tracking[col] || self.dict_tracking[col] !=
DictTracking::Disabled;
+
+ for row in nulls.valid_indices() {
+ let before = self.value_buffers[col].len();
+
self.value_buffers[col].extend_from_slice(&values[row].to_be_bytes());
+ self.non_null_counts[col] += 1;
+
+ if track_encoding {
+ self.track_encoding_value(col, before, 8);
+ track_encoding =
+ self.const_tracking[col] || self.dict_tracking[col] !=
DictTracking::Disabled;
+ }
+ }
+
+ valid_count * 8
+ }
+
+ fn track_encoding_value(&mut self, col: usize, value_start: usize,
value_len: usize) {
+ if self.non_null_counts[col] == 1 {
+ self.first_value_len[col] = value_len;
+ return;
+ }
+
+ if self.const_tracking[col] {
+ if value_len == self.first_value_len[col]
+ && equals_first_value(&self.value_buffers[col], value_start,
value_len)
+ {
+ return;
+ }
+
+ self.const_tracking[col] = false;
+ self.activate_dict_tracking(col, value_start, value_len);
+ } else {
+ self.track_dict_value(col, value_start, value_len);
+ }
+ }
+
+ fn activate_dict_tracking(&mut self, col: usize, value_start: usize,
value_len: usize) {
+ debug_assert_eq!(self.dict_tracking[col], DictTracking::Pending);
+ self.dict_total_bytes[col] = 0;
+ if uses_long_dict(self.fixed_widths[col]) {
+ self.long_dict_maps[col]
+ .get_or_insert_with(HashMap::new)
+ .clear();
+ } else {
+ self.byte_dict_maps[col]
+ .get_or_insert_with(HashMap::new)
+ .clear();
+ }
+ self.dict_tracking[col] = DictTracking::Active;
+
+ self.track_dict_value(col, 0, self.first_value_len[col]);
+ if self.dict_tracking[col] == DictTracking::Active {
+ self.track_dict_value(col, value_start, value_len);
+ }
+ }
+
+ fn track_dict_value(&mut self, col: usize, value_start: usize, value_len:
usize) {
+ if self.dict_tracking[col] != DictTracking::Active {
+ return;
+ }
+
+ let disable = if let Some(ref mut dict) = self.long_dict_maps[col] {
+ let key = values::extract_fixed_key(
+ &self.value_buffers[col],
+ value_start,
+ self.fixed_widths[col],
+ );
+ let len = dict.len();
+ dict.entry(key).or_insert(len);
+ dict.len() > self.max_dict_entries
+ } else if let Some(ref mut dict) = self.byte_dict_maps[col] {
+ let value = &self.value_buffers[col][value_start..value_start +
value_len];
+ if !dict.contains_key(value) {
+ let len = dict.len();
+ dict.insert(value.to_vec(), len);
+ self.dict_total_bytes[col] += value_len;
+ }
+ dict.len() > self.max_dict_entries
+ || self.dict_total_bytes[col] > self.max_dict_total_bytes
+ } else {
+ unreachable!("active dictionary tracking must have a dictionary
map")
+ };
+
+ if disable {
+ self.dict_tracking[col] = DictTracking::Disabled;
+ self.long_dict_maps[col] = None;
+ self.byte_dict_maps[col] = None;
+ }
+ }
+
fn compute_encodings(&self) -> (Vec<u8>, Vec<bool>) {
let mut encodings = vec![0u8; self.total_columns];
let mut has_nulls = vec![false; self.total_columns];
@@ -760,24 +814,19 @@ impl BucketWriter {
pub fn reset(&mut self) {
for i in 0..self.total_columns {
- for b in &mut self.null_bitmaps[i] {
- *b = 0;
- }
+ self.null_bitmaps[i].fill(0);
self.value_buffers[i].clear();
self.non_null_counts[i] = 0;
self.const_tracking[i] = true;
self.first_value_len[i] = 0;
+ self.dict_tracking[i] = DictTracking::Pending;
self.dict_total_bytes[i] = 0;
if uses_long_dict(self.fixed_widths[i]) {
if let Some(ref mut dict) = self.long_dict_maps[i] {
dict.clear();
- } else {
- self.long_dict_maps[i] = Some(HashMap::new());
}
} else if let Some(ref mut dict) = self.byte_dict_maps[i] {
dict.clear();
- } else {
- self.byte_dict_maps[i] = Some(HashMap::new());
}
}
self.num_rows = 0;
@@ -787,6 +836,9 @@ impl BucketWriter {
}
fn get_dict_size(&self, col: usize) -> usize {
+ if self.dict_tracking[col] != DictTracking::Active {
+ return 0;
+ }
if let Some(ref dict) = self.long_dict_maps[col] {
return dict.len();
}
@@ -833,6 +885,9 @@ impl BucketWriter {
}
fn dict_encoded_size(&self, col: usize) -> usize {
+ if self.dict_tracking[col] != DictTracking::Active {
+ return usize::MAX;
+ }
let (num_entries, entry_bytes) = if let Some(ref dict) =
self.long_dict_maps[col] {
(dict.len(), dict.len() * self.fixed_widths[col] as usize)
} else if let Some(ref dict) = self.byte_dict_maps[col] {
@@ -891,6 +946,41 @@ impl BucketWriter {
}
}
+fn append_null_bitmap(bitmap: &mut [u8], start_row: usize, nulls: &NullBuffer)
{
+ debug_assert!(start_row + nulls.len() <= bitmap.len() * 8);
+
+ let chunks = nulls.inner().bit_chunks();
+ for (chunk, valid_bits) in chunks.iter().enumerate() {
+ or_null_bits(bitmap, start_row + chunk * 64, !valid_bits, 64);
+ }
+
+ let remainder_len = nulls.len() % 64;
+ if remainder_len != 0 {
+ let remainder_mask = (1u64 << remainder_len) - 1;
+ let null_bits = !chunks.remainder_bits() & remainder_mask;
+ or_null_bits(
+ bitmap,
+ start_row + nulls.len() - remainder_len,
+ null_bits,
+ remainder_len,
+ );
+ }
+}
+
+#[inline]
+fn or_null_bits(bitmap: &mut [u8], mut bit_offset: usize, mut bits: u64, mut
len: usize) {
+ while len != 0 {
+ let byte_offset = bit_offset / 8;
+ let bit_in_byte = bit_offset % 8;
+ let take = len.min(8 - bit_in_byte);
+ let mask = ((1u16 << take) - 1) as u8;
+ bitmap[byte_offset] |= ((bits as u8) & mask) << bit_in_byte;
+ bits >>= take;
+ bit_offset += take;
+ len -= take;
+ }
+}
+
enum TypedArrayRef<'a> {
Boolean(&'a BooleanArray),
Int8(&'a Int8Array),
@@ -1730,6 +1820,211 @@ mod tests {
assert_eq!(data[h] & 0x03, ENCODING_CONST);
}
+ #[test]
+ fn test_const_columns_do_not_build_dictionaries() {
+ let types = [DataType::Int32, DataType::Utf8];
+ let type_refs: Vec<&DataType> = types.iter().collect();
+ let mut writer = BucketWriter::new(&type_refs, 32768, 255);
+
+ let ints = Int32Array::from(vec![42; 100]);
+ let strings = StringArray::from(vec!["same"; 100]);
+ writer
+ .write_columns(&[&ints, &strings], &[&types[0], &types[1]])
+ .unwrap();
+
+ assert!(writer.long_dict_maps[0].is_none());
+ assert!(writer.byte_dict_maps[1].is_none());
+ assert_eq!(writer.dict_total_bytes[1], 0);
+ }
+
+ #[test]
+ fn test_dictionary_tracking_starts_on_first_non_const_value() {
+ let types = [DataType::Int32, DataType::Utf8];
+ let type_refs: Vec<&DataType> = types.iter().collect();
+ let mut writer = BucketWriter::new(&type_refs, 32768, 255);
+
+ let const_ints = Int32Array::from(vec![42; 50]);
+ let const_strings = StringArray::from(vec!["same"; 50]);
+ writer
+ .write_columns(&[&const_ints, &const_strings], &[&types[0],
&types[1]])
+ .unwrap();
+
+ let varied_ints = Int32Array::from(vec![42, 7, 42, 7]);
+ let varied_strings = StringArray::from(vec!["same", "other", "same",
"other"]);
+ writer
+ .write_columns(&[&varied_ints, &varied_strings], &[&types[0],
&types[1]])
+ .unwrap();
+
+ assert_eq!(writer.dict_tracking, vec![DictTracking::Active; 2]);
+ assert_eq!(writer.long_dict_maps[0].as_ref().unwrap().len(), 2);
+ assert_eq!(writer.byte_dict_maps[1].as_ref().unwrap().len(), 2);
+
+ let data = writer.finish();
+ assert_eq!(data[0] & 0x03, ENCODING_DICT);
+ assert_eq!((data[0] >> 2) & 0x03, ENCODING_DICT);
+ }
+
+ #[test]
+ fn test_dictionary_tracking_disables_after_entry_limit() {
+ let types = [DataType::Int32, DataType::Utf8];
+ let type_refs: Vec<&DataType> = types.iter().collect();
+ let mut writer = BucketWriter::new(&type_refs, 32768, 255);
+
+ let ints = Int32Array::from_iter_values((0..255).chain(0..255));
+ let string_values: Vec<String> = (0..255)
+ .chain(0..255)
+ .map(|i| format!("value-{i:03}"))
+ .collect();
+ let strings =
StringArray::from_iter_values(string_values.iter().map(String::as_str));
+ writer
+ .write_columns(&[&ints, &strings], &[&types[0], &types[1]])
+ .unwrap();
+
+ assert_eq!(writer.dict_tracking, vec![DictTracking::Active; 2]);
+ assert_eq!(writer.long_dict_maps[0].as_ref().unwrap().len(), 255);
+ assert_eq!(writer.byte_dict_maps[1].as_ref().unwrap().len(), 255);
+ let data = writer.finish();
+ assert_eq!(data[0] & 0x03, ENCODING_DICT);
+ assert_eq!((data[0] >> 2) & 0x03, ENCODING_DICT);
+
+ let next_int = Int32Array::from(vec![255]);
+ let next_string = StringArray::from(vec!["value-255"]);
+ writer
+ .write_columns(&[&next_int, &next_string], &[&types[0], &types[1]])
+ .unwrap();
+
+ assert_eq!(writer.dict_tracking, vec![DictTracking::Disabled; 2]);
+ assert!(writer.long_dict_maps.iter().all(Option::is_none));
+ assert!(writer.byte_dict_maps.iter().all(Option::is_none));
+
+ let later_ints = Int32Array::from(vec![0, 256]);
+ let later_strings = StringArray::from(vec!["value-000", "value-256"]);
+ writer
+ .write_columns(&[&later_ints, &later_strings], &[&types[0],
&types[1]])
+ .unwrap();
+
+ assert_eq!(writer.dict_tracking, vec![DictTracking::Disabled; 2]);
+ assert!(writer.long_dict_maps.iter().all(Option::is_none));
+ assert!(writer.byte_dict_maps.iter().all(Option::is_none));
+ let data = writer.finish();
+ assert_eq!(data[0] & 0x03, ENCODING_PLAIN);
+ assert_eq!((data[0] >> 2) & 0x03, ENCODING_PLAIN);
+ }
+
+ #[test]
+ fn test_dictionary_tracking_disables_after_byte_limit() {
+ let types = [DataType::Utf8];
+ let type_refs: Vec<&DataType> = types.iter().collect();
+ // Encoded dictionary keys include the one-byte string length prefix:
+ // "a" is 2 bytes and "bb" is 3 bytes, exactly filling this budget.
+ let mut writer = BucketWriter::new(&type_refs, 5, 255);
+
+ // Repeat enough rows so DICT is smaller than PLAIN and the encoding
assertion is stable.
+ let mut within_budget: Vec<Option<&str>> = (0..64)
+ .map(|i| Some(if i % 2 == 0 { "a" } else { "bb" }))
+ .collect();
+ within_budget.push(None);
+ let within_budget = StringArray::from(within_budget);
+ writer
+ .write_columns(&[&within_budget], &[&DataType::Utf8])
+ .unwrap();
+
+ assert_eq!(writer.dict_tracking[0], DictTracking::Active);
+ assert_eq!(writer.byte_dict_maps[0].as_ref().unwrap().len(), 2);
+ assert_eq!(writer.dict_total_bytes[0], 5);
+ let data = writer.finish();
+ assert_eq!(data[0] & 0x03, ENCODING_DICT);
+
+ let beyond_budget = StringArray::from(vec![Some("a"), Some("ccc"),
None]);
+ writer
+ .write_columns(&[&beyond_budget], &[&DataType::Utf8])
+ .unwrap();
+
+ assert_eq!(writer.dict_tracking[0], DictTracking::Disabled);
+ assert!(writer.byte_dict_maps[0].is_none());
+ let data = writer.finish();
+ assert_eq!(data[0] & 0x03, ENCODING_PLAIN);
+ }
+
+ #[test]
+ fn test_append_null_bitmap_handles_source_and_destination_offsets() {
+ let validity = BooleanBuffer::from(vec![
+ true, false, true, false, false, true, true, false, true, false,
true,
+ ]);
+ let nulls = NullBuffer::new(validity).slice(2, 7);
+ let mut bitmap = vec![0b0000_0101, 0, 0];
+
+ append_null_bitmap(&mut bitmap, 3, &nulls);
+
+ assert_eq!(bitmap, vec![0b0011_0101, 0b0000_0001, 0]);
+ }
+
+ #[test]
+ fn test_append_null_bitmap_matches_scalar_copy_across_chunk_boundaries() {
+ for source_offset in 0..8 {
+ for destination_offset in 0..16 {
+ for len in [1, 7, 8, 9, 63, 64, 65, 127, 128, 129] {
+ let validity: Vec<bool> = (0..source_offset + len)
+ .map(|i| i % 3 != 0 && i % 11 != 0)
+ .collect();
+ let nulls =
NullBuffer::new(BooleanBuffer::from(validity.clone()))
+ .slice(source_offset, len);
+ let mut actual = vec![0u8; (destination_offset +
len).div_ceil(8) + 1];
+ let mut expected = actual.clone();
+
+ for row in 0..len {
+ if !validity[source_offset + row] {
+ let bit = destination_offset + row;
+ expected[bit / 8] |= 1 << (bit % 8);
+ }
+ }
+ append_null_bitmap(&mut actual, destination_offset,
&nulls);
+
+ assert_eq!(
+ actual, expected,
+ "source_offset={source_offset},
destination_offset={destination_offset}, len={len}"
+ );
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn
test_nullable_timestamp_millis_batch_preserves_slices_and_append_offsets() {
+ let data_type =
DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None);
+ let mut writer = BucketWriter::new(&[&data_type], 32768, 255);
+
+ let source = TimestampMillisecondArray::from(vec![
+ Some(99),
+ Some(10),
+ None,
+ Some(10),
+ Some(20),
+ None,
+ ]);
+ let first = source.slice(1, 4);
+ let first_size = writer.write_columns(&[&first],
&[&data_type]).unwrap();
+ assert_eq!(first_size, 28);
+
+ let second = TimestampMillisecondArray::from(vec![None, Some(20),
None]);
+ let second_size = writer.write_columns(&[&second],
&[&data_type]).unwrap();
+ assert_eq!(second_size, 11);
+
+ assert_eq!(writer.num_rows, 7);
+ assert_eq!(writer.non_null_counts[0], 4);
+ assert_eq!(writer.null_bitmaps[0][0], 0b0101_0010);
+ assert_eq!(
+ writer.value_buffers[0],
+ [10_i64, 10, 20, 20]
+ .into_iter()
+ .flat_map(i64::to_be_bytes)
+ .collect::<Vec<_>>()
+ );
+ assert!(!writer.const_tracking[0]);
+ assert_eq!(writer.dict_tracking[0], DictTracking::Active);
+ assert_eq!(writer.long_dict_maps[0].as_ref().unwrap().len(), 2);
+ }
+
#[test]
fn test_dict_string_encoding() {
let types = [DataType::Utf8];
