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 3 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
157 changes: 74 additions & 83 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 @@ -4134,14 +4134,12 @@ impl PrimitiveStructuralEncoder {
/// 1. Dictionary: stores unique values
/// 2. Indices: maps each value to a dictionary entry
///
/// Indices: num_values × 4 bytes (32-bit i32)
/// For FixedWidth (e.g., 128-bit Decimal):
/// - Dictionary: cardinality × 16 bytes (128 bits per value)
/// - Indices: num_values × 4 bytes (32-bit i32)
/// - Dictionary values: cardinality × 16 bytes (128 bits per value)
///
/// For VariableWidth (strings/binary):
/// - Dictionary values: cardinality × avg_value_size (actual data)
/// - Dictionary offsets: cardinality × offset_size (32 or 64 bits)
/// - Indices: num_values × offset_size (same as dictionary offsets)
fn estimate_dict_size(data_block: &DataBlock) -> Option<u64> {
let cardinality = if let Some(cardinality_array) = data_block.get_stat(Stat::Cardinality) {
cardinality_array.as_primitive::<UInt64Type>().value(0)
Expand All @@ -4150,36 +4148,15 @@ impl PrimitiveStructuralEncoder {
};

let num_values = data_block.num_values();
let indices_size = num_values * (DICT_INDICES_BITS_PER_VALUE / 8);

match data_block {
DataBlock::FixedWidth(_) => {
// Dictionary: cardinality unique values at 128 bits each
let dict_size = cardinality * (DICT_FIXED_WIDTH_BITS_PER_VALUE / 8);
// Indices: num_values indices at 32 bits each
let indices_size = num_values * (DICT_INDICES_BITS_PER_VALUE / 8);
Some(dict_size + indices_size)
}
DataBlock::VariableWidth(var) => {
// Only 32-bit and 64-bit offsets are supported
if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
return None;
}
let bits_per_offset = var.bits_per_offset as u64;

let data_size = data_block.data_size();
let avg_value_size = data_size / 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
let indices_size = num_values * (bits_per_offset / 8);
let avg_value_size = match data_block {
DataBlock::FixedWidth(_) => DICT_FIXED_WIDTH_BITS_PER_VALUE / 8,
DataBlock::VariableWidth(_) => data_block.data_size() / num_values,
_ => unreachable!("Variable width offsets can only be 32 or 64 bits"),
};

Some(dict_values_size + dict_offsets_size + indices_size)
}
_ => None,
}
Some(avg_value_size * cardinality + indices_size)
}

fn should_dictionary_encode(data_block: &DataBlock, field: &Field) -> bool {
Expand All @@ -4192,6 +4169,13 @@ impl PrimitiveStructuralEncoder {
return false;
}

// Currently VariableWidth only supports 32 and 64 bits
if let DataBlock::VariableWidth(var) = data_block {
if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
return false;
}
}

// Don't dictionary encode tiny arrays
let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
.ok()
Expand Down Expand Up @@ -4341,57 +4325,64 @@ impl PrimitiveStructuralEncoder {
Some(dictionary_data_block),
num_rows
)
} else if Self::should_dictionary_encode(&data_block, &field) {
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,
repdefs,
row_number,
Some(dictionary_data_block),
num_rows,
)
} 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,
repdefs,
row_number,
None,
num_rows,
)
} 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,
repdefs,
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
let dict_result = if Self::should_dictionary_encode(&data_block, &field) {
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,
repdefs,
row_number,
Some(dictionary_data_block),
num_rows,
)
} 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,
repdefs,
row_number,
None,
num_rows,
)
} 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,
repdefs,
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 @@ -5611,7 +5602,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