mirror of
https://github.com/quickwit-oss/tantivy.git
synced 2026-10-06 11:52:40 +00:00
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:
@@ -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()))
|
||||
|
||||
@@ -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)));
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user