Skip to content

feat: abort dictionary encode if not useful #5055

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Xuanwo merged 13 commits into from
Jan 28, 2026
Merged

feat: abort dictionary encode if not useful #5055

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
8b4eb37
feat: fix 64 bit variable block to use 32 bits index
Oct 20, 2025
641b8b2
fix clippy
Oct 23, 2025
b7e193a
abort dictionary encoding on flight if it's not useful as the estimat…
Oct 23, 2025
7fa19c4
Merge remote-tracking branch 'origin/main' into abortDictionaryEncode…
Xuanwo Jan 22, 2026
6dac596
fix: avoid non-final 1-value miniblock chunks
Xuanwo Jan 22, 2026
f11ac05
perf: abort dictionary encode when not beneficial
Xuanwo Jan 22, 2026
f3ca17b
Merge branch 'main' into abortDictionaryEncodeIfNotUseful
Xuanwo Jan 27, 2026
0eaa213
Fix
Xuanwo Jan 28, 2026
a1b9c03
Merge remote-tracking branch 'origin/main' into abortDictionaryEncode…
Xuanwo Jan 28, 2026
ba52fc7
Merge remote-tracking branch 'origin/main' into abortDictionaryEncode…
Xuanwo Jan 28, 2026
fef8666
fix: resolve origin/main merge conflicts
Xuanwo Jan 28, 2026
731cd33
Merge remote-tracking branch 'pr-head/abortDictionaryEncodeIfNotUsefu…
Xuanwo Jan 28, 2026
793202f
Merge remote-tracking branch 'origin/main' into abortDictionaryEncode…
Xuanwo Jan 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 69 additions & 67 deletions rust/lance-encoding/src/encodings/logical/primitive.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4477,7 +4477,6 @@ impl PrimitiveStructuralEncoder {
///
/// For VariableWidth (strings/binary):
/// - Dictionary values: cardinality × avg_value_size (actual data)
/// - Dictionary offsets: cardinality × offset_size (32 or 64 bits)
/// - Indices: num_values × 4 bytes (32-bit i32)
fn estimate_dict_size(data_block: &DataBlock, version: LanceFileVersion) -> Option<u64> {
let cardinality = if let Some(cardinality_array) = data_block.get_stat(Stat::Cardinality) {
Expand All @@ -4490,7 +4489,6 @@ impl PrimitiveStructuralEncoder {
if num_values == 0 {
return None;
}

match data_block {
DataBlock::FixedWidth(fixed) => {
if fixed.bits_per_value == 64 && version < LanceFileVersion::V2_2 {
Expand Down Expand Up @@ -4518,23 +4516,20 @@ impl PrimitiveStructuralEncoder {
if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
return None;
}
let bits_per_offset = var.bits_per_offset as u64;
// Dictionary indices are always i32.
if cardinality > i32::MAX as u64 {
return None;
}

let data_size = data_block.data_size();
let avg_value_size = data_size / num_values;
let bytes_per_offset = var.bits_per_offset as u64 / 8;
let avg_value_size = (var.data.len() as u64) / num_values;

// Dictionary values: actual bytes of unique strings/binary
let dict_values_size = cardinality * avg_value_size;
// Dictionary offsets: pointers into dictionary values
let dict_offsets_size = cardinality * (bits_per_offset / 8);
// Indices: map each row to dictionary entry (always i32)
let indices_size = num_values * (DICT_INDICES_BITS_PER_VALUE / 8);
let dict_values_size = cardinality.checked_mul(avg_value_size)?;
let dict_offsets_size = cardinality.checked_mul(bytes_per_offset)?;
let indices_size = num_values.checked_mul(DICT_INDICES_BITS_PER_VALUE / 8)?;

Some(dict_values_size + dict_offsets_size + indices_size)
dict_values_size
.checked_add(dict_offsets_size)?
.checked_add(indices_size)
}
_ => None,
}
Expand Down Expand Up @@ -4741,60 +4736,67 @@ impl PrimitiveStructuralEncoder {
num_rows,
support_large_chunk,
)
} else if Self::should_dictionary_encode(&data_block, &field, version) {
log::debug!(
"Encoding column {} with {} items using dictionary encoding (mini-block layout)",
column_idx,
num_values
);
let (indices_data_block, dictionary_data_block) =
dict::dictionary_encode(data_block);

Self::encode_miniblock(
column_idx,
&field,
compression_strategy.as_ref(),
indices_data_block,
repdef,
row_number,
Some(dictionary_data_block),
num_rows,
support_large_chunk,
)
} else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
log::debug!(
"Encoding column {} with {} items using mini-block layout",
column_idx,
num_values
);
Self::encode_miniblock(
column_idx,
&field,
compression_strategy.as_ref(),
data_block,
repdef,
row_number,
None,
num_rows,
support_large_chunk,
)
} else if Self::prefers_fullzip(encoding_metadata.as_ref()) {
log::debug!(
"Encoding column {} with {} items using full-zip layout",
column_idx,
num_values
);
Self::encode_full_zip(
column_idx,
&field,
compression_strategy.as_ref(),
data_block,
repdef,
row_number,
num_rows,
)
} else {
Err(Error::InvalidInput { source: format!("Cannot determine structural encoding for field {}. This typically indicates an invalid value of the field metadata key {}", field.name, STRUCTURAL_ENCODING_META_KEY).into(), location: location!() })
// Try dictionary encoding first if applicable. If encoding aborts, fall back to the
// preferred structural encoding.
let dict_result = if Self::should_dictionary_encode(&data_block, &field, version) {
log::debug!(
"Encoding column {} with {} items using dictionary encoding (mini-block layout)",
column_idx,
num_values
);
dict::dictionary_encode(data_block.clone())
} else {
None
};

if let Some((indices_data_block, dictionary_data_block)) = dict_result {
Self::encode_miniblock(
column_idx,
&field,
compression_strategy.as_ref(),
indices_data_block,
repdef,
row_number,
Some(dictionary_data_block),
num_rows,
support_large_chunk,
)
} else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
log::debug!(
"Encoding column {} with {} items using mini-block layout",
column_idx,
num_values
);
Self::encode_miniblock(
column_idx,
&field,
compression_strategy.as_ref(),
data_block,
repdef,
row_number,
None,
num_rows,
support_large_chunk,
)
} else if Self::prefers_fullzip(encoding_metadata.as_ref()) {
log::debug!(
"Encoding column {} with {} items using full-zip layout",
column_idx,
num_values
);
Self::encode_full_zip(
column_idx,
&field,
compression_strategy.as_ref(),
data_block,
repdef,
row_number,
num_rows,
)
} else {
Err(Error::InvalidInput { source: format!("Cannot determine structural encoding for field {}. This typically indicates an invalid value of the field metadata key {}", field.name, STRUCTURAL_ENCODING_META_KEY).into(), location: location!() })
}
}
})
.boxed();
Expand Down Expand Up @@ -6157,7 +6159,7 @@ mod tests {
let data_size = block.data_size();
let avg_value_size = data_size / 1000;

let expected = 400 * avg_value_size + 400 * 4 + 1000 * 4;
let expected = 400 * avg_value_size + 1000 * 4;

assert_eq!(estimated_size, expected);
}
Expand Down
Loading