Classify tie breakers during columnar merge

Represent generated tie breakers as a merge column category so existing
values are skipped and regenerated without being loaded. Test grouping,
block linearity, and compact merged output.
This commit is contained in:
Pascal Seitz
2026-09-15 09:45:58 +02:00
parent e7ba845037
commit a4cde3a22c
2 changed files with 90 additions and 13 deletions
+23 -5
View File
@@ -30,9 +30,10 @@ use crate::{
///
/// See also [README.md].
///
/// The ordering has to match the ordering of the variants in [ColumnType].
/// Except for generated tie-breakers, the ordering matches the variants in [ColumnType].
#[derive(Copy, Clone, Eq, PartialOrd, Ord, PartialEq, Hash, Debug)]
pub(crate) enum ColumnTypeCategory {
TieBreaker,
Numerical,
Bytes,
Str,
@@ -88,10 +89,11 @@ pub fn merge_columnar(
.map(|reader| reader.num_docs())
.collect::<Vec<u32>>();
let columns_to_merge = group_columns_for_merge(columnar_readers, required_columns)?;
let columns_to_merge =
group_columns_for_merge(columnar_readers, required_columns, tie_breaker_columns)?;
for res in columns_to_merge {
let ((column_name, _column_type_category), grouped_columns) = res;
if tie_breaker_columns.iter().any(|name| name == &column_name) {
let ((column_name, column_type_category), grouped_columns) = res;
if column_type_category == ColumnTypeCategory::TieBreaker {
let mut column_serializer =
serializer.start_serialize_column(column_name.as_bytes(), ColumnType::U64);
crate::column::serialize_generated_tie_breaker_column(
@@ -411,21 +413,37 @@ fn is_empty_after_merge(
fn group_columns_for_merge<'a>(
columnar_readers: &'a [&'a ColumnarReader],
required_columns: &'a [(String, ColumnType)],
tie_breaker_columns: &[String],
) -> io::Result<BTreeMap<(String, ColumnTypeCategory), GroupedColumnsHandle>> {
let mut columns: BTreeMap<(String, ColumnTypeCategory), GroupedColumnsHandle> = BTreeMap::new();
let tie_breaker_columns: HashSet<&str> =
tie_breaker_columns.iter().map(String::as_str).collect();
for &(ref column_name, column_type) in required_columns {
if tie_breaker_columns.contains(column_name.as_str()) {
continue;
}
columns
.entry((column_name.clone(), column_type.into()))
.or_insert_with(|| GroupedColumnsHandle::new(columnar_readers.len()))
.require_type(column_type)?;
}
for column_name in &tie_breaker_columns {
columns
.entry(((*column_name).to_string(), ColumnTypeCategory::TieBreaker))
.or_insert_with(|| GroupedColumnsHandle::new(columnar_readers.len()))
.require_type(ColumnType::U64)?;
}
for (columnar_id, columnar_reader) in columnar_readers.iter().enumerate() {
let column_name_and_handle = columnar_reader.iter_columns()?;
for (column_name, handle) in column_name_and_handle {
let column_category: ColumnTypeCategory = handle.column_type().into();
if tie_breaker_columns.contains(column_name.as_str()) {
continue;
}
let column_category = handle.column_type().into();
columns
.entry((column_name, column_category))
.or_insert_with(|| GroupedColumnsHandle::new(columnar_readers.len()))
+67 -8
View File
@@ -30,7 +30,7 @@ fn test_column_coercion_to_u64() {
let columnar2 = make_columnar("numbers", &[u64::MAX]);
let columnars = &[&columnar1, &columnar2];
let column_map: BTreeMap<(String, ColumnTypeCategory), GroupedColumnsHandle> =
group_columns_for_merge(columnars, &[]).unwrap();
group_columns_for_merge(columnars, &[], &[]).unwrap();
assert_eq!(column_map.len(), 1);
assert!(column_map.contains_key(&("numbers".to_string(), ColumnTypeCategory::Numerical)));
}
@@ -41,7 +41,7 @@ fn test_column_coercion_to_i64() {
let columnar2 = make_columnar("numbers", &[2u64]);
let columnars = &[&columnar1, &columnar2];
let column_map: BTreeMap<(String, ColumnTypeCategory), GroupedColumnsHandle> =
group_columns_for_merge(columnars, &[]).unwrap();
group_columns_for_merge(columnars, &[], &[]).unwrap();
assert_eq!(column_map.len(), 1);
assert!(column_map.contains_key(&("numbers".to_string(), ColumnTypeCategory::Numerical)));
}
@@ -65,19 +65,77 @@ fn test_group_columns_with_required_column() {
let columnar2 = make_columnar("numbers", &[2u64]);
let columnars = &[&columnar1, &columnar2];
let column_map: BTreeMap<(String, ColumnTypeCategory), GroupedColumnsHandle> =
group_columns_for_merge(columnars, &[("numbers".to_string(), ColumnType::U64)]).unwrap();
group_columns_for_merge(columnars, &[("numbers".to_string(), ColumnType::U64)], &[])
.unwrap();
assert_eq!(column_map.len(), 1);
assert!(column_map.contains_key(&("numbers".to_string(), ColumnTypeCategory::Numerical)));
}
#[test]
fn test_group_tie_breaker_column() {
let columnar = make_columnar("tie", &[1u64]);
let column_map = group_columns_for_merge(&[&columnar], &[], &["tie".to_string()]).unwrap();
assert_eq!(column_map.len(), 1);
let grouped = column_map
.get(&("tie".to_string(), ColumnTypeCategory::TieBreaker))
.unwrap();
assert!(grouped.columns[0].is_none());
}
#[test]
fn test_merge_generated_tie_breaker_columns_stays_small() {
fn make_tie_breaker_columnar(num_docs: u32) -> ColumnarReader {
let mut writer = ColumnarWriter::default();
writer.record_tie_breaker_column("tie");
let mut buffer = Vec::new();
writer.serialize(num_docs, None, &mut buffer).unwrap();
ColumnarReader::open(buffer).unwrap()
}
let columnar1 = make_tie_breaker_columnar(10_000);
let columnar2 = make_tie_breaker_columnar(10_000);
let columnars = &[&columnar1, &columnar2];
let mut buffer = Vec::new();
merge_columnar(
columnars,
&[("tie".to_string(), ColumnType::U64)],
&["tie".to_string()],
StackMergeOrder::stack(columnars).into(),
&mut buffer,
)
.unwrap();
assert!(
buffer.len() <= 600,
"merged columnar is {} bytes",
buffer.len()
);
let merged = ColumnarReader::open(buffer).unwrap();
let column = merged.read_columns("tie").unwrap()[0]
.open_u64_lenient()
.unwrap()
.unwrap();
assert_eq!(merged.num_docs(), 20_000);
assert_eq!(column.get_cardinality(), Cardinality::Full);
for block_start in (0..20_000).step_by(512) {
let block_end = (block_start + 512).min(20_000);
for doc in block_start + 1..block_end {
assert_eq!(column.first(doc), Some(column.first(doc - 1).unwrap() + 1));
}
}
}
#[test]
fn test_group_columns_required_column_with_no_existing_columns() {
let columnar1 = make_columnar("numbers", &[2u64]);
let columnar2 = make_columnar("numbers", &[2u64]);
let columnars = &[&columnar1, &columnar2];
let column_map: BTreeMap<_, _> =
group_columns_for_merge(columnars, &[("required_col".to_string(), ColumnType::Str)])
.unwrap();
let column_map: BTreeMap<_, _> = group_columns_for_merge(
columnars,
&[("required_col".to_string(), ColumnType::Str)],
&[],
)
.unwrap();
assert_eq!(column_map.len(), 2);
let columns = &column_map
.get(&("required_col".to_string(), ColumnTypeCategory::Str))
@@ -94,7 +152,8 @@ fn test_group_columns_required_column_is_above_all_columns_have_the_same_type_ru
let columnar2 = make_columnar("numbers", &[2i64]);
let columnars = &[&columnar1, &columnar2];
let column_map: BTreeMap<(String, ColumnTypeCategory), GroupedColumnsHandle> =
group_columns_for_merge(columnars, &[("numbers".to_string(), ColumnType::U64)]).unwrap();
group_columns_for_merge(columnars, &[("numbers".to_string(), ColumnType::U64)], &[])
.unwrap();
assert_eq!(column_map.len(), 1);
assert!(column_map.contains_key(&("numbers".to_string(), ColumnTypeCategory::Numerical)));
}
@@ -105,7 +164,7 @@ fn test_missing_column() {
let columnar2 = make_columnar("numbers2", &[2u64]);
let columnars = &[&columnar1, &columnar2];
let column_map: BTreeMap<(String, ColumnTypeCategory), GroupedColumnsHandle> =
group_columns_for_merge(columnars, &[]).unwrap();
group_columns_for_merge(columnars, &[], &[]).unwrap();
assert_eq!(column_map.len(), 2);
assert!(column_map.contains_key(&("numbers".to_string(), ColumnTypeCategory::Numerical)));
{