From 46ee982b520da71a1cadce588eca8244140cf394 Mon Sep 17 00:00:00 2001 From: Pascal Seitz Date: Mon, 14 Sep 2026 17:26:01 +0200 Subject: [PATCH] add tie_breaker option --- columnar/Cargo.toml | 2 +- columnar/src/column/mod.rs | 1 + columnar/src/column/serialize.rs | 17 +++++++ columnar/src/column_values/mod.rs | 1 + .../u64_based/blockwise_linear.rs | 2 +- columnar/src/column_values/u64_based/mod.rs | 2 +- columnar/src/columnar/merge/mod.rs | 28 +++++++++++ columnar/src/columnar/mod.rs | 2 +- columnar/src/columnar/reader/mod.rs | 19 ++++++++ columnar/src/columnar/writer/mod.rs | 48 ++++++++++++------- columnar/src/lib.rs | 2 +- src/fastfield/mod.rs | 22 +++++++++ src/fastfield/plugin.rs | 18 ++++++- src/fastfield/writer.rs | 7 +++ src/schema/numeric_options.rs | 25 ++++++++++ src/schema/schema.rs | 12 +++++ 16 files changed, 185 insertions(+), 23 deletions(-) diff --git a/columnar/Cargo.toml b/columnar/Cargo.toml index 10b49a24d..6e07375c5 100644 --- a/columnar/Cargo.toml +++ b/columnar/Cargo.toml @@ -18,11 +18,11 @@ common = { version= "0.11", path = "../common", package = "tantivy-common" } tantivy-bitpacker = { version= "0.10", path = "../bitpacker/" } serde = "1.0.152" downcast-rs = "2.0.1" +rand = "0.9" [dev-dependencies] proptest = "1" more-asserts = "0.3.1" -rand = "0.9" binggan = "0.17.0" [[bench]] diff --git a/columnar/src/column/mod.rs b/columnar/src/column/mod.rs index f6a50b45f..3bc61cba0 100644 --- a/columnar/src/column/mod.rs +++ b/columnar/src/column/mod.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use common::BinarySerializable; pub use dictionary_encoded::{BytesColumn, StrColumn}; +pub(crate) use serialize::serialize_generated_tie_breaker_column; pub use serialize::{ open_column_bytes, open_column_str, open_column_u64, open_column_u128, open_column_u128_as_compact_u64, serialize_column_mappable_to_u64, diff --git a/columnar/src/column/serialize.rs b/columnar/src/column/serialize.rs index e2127933b..75fa52f9a 100644 --- a/columnar/src/column/serialize.rs +++ b/columnar/src/column/serialize.rs @@ -3,6 +3,7 @@ use std::io::Write; use std::sync::Arc; use common::OwnedBytes; +use rand::Rng; use sstable::Dictionary; use crate::column::{BytesColumn, Column}; @@ -25,6 +26,22 @@ pub fn serialize_column_mappable_to_u128( Ok(()) } +pub(crate) fn serialize_generated_tie_breaker_column( + num_docs: u32, + output: &mut impl Write, +) -> io::Result<()> { + let block_size = crate::column_values::BLOCK_SIZE; + let max_start = u32::MAX - (block_size - 1); + let mut rng = rand::rng(); + let mut values = Vec::with_capacity(num_docs as usize); + for block_start_doc in (0..num_docs).step_by(block_size as usize) { + let start = rng.random_range(0..=max_start); + let block_len = (num_docs - block_start_doc).min(block_size); + values.extend((0..block_len).map(|offset| (start + offset) as u64)); + } + serialize_column_mappable_to_u64(SerializableColumnIndex::Full, &&values[..], output) +} + pub fn serialize_column_mappable_to_u64( column_index: SerializableColumnIndex<'_>, column_values: &impl Iterable, diff --git a/columnar/src/column_values/mod.rs b/columnar/src/column_values/mod.rs index 64bc69b25..0cb4ac1db 100644 --- a/columnar/src/column_values/mod.rs +++ b/columnar/src/column_values/mod.rs @@ -26,6 +26,7 @@ mod monotonic_column; pub(crate) use merge::MergedColumnValues; pub use stats::ColumnStats; +pub(crate) use u64_based::blockwise_linear::BLOCK_SIZE; pub use u64_based::{ ALL_U64_CODEC_TYPES, CodecType, load_u64_based_column_values, serialize_and_load_u64_based_column_values, serialize_u64_based_column_values, diff --git a/columnar/src/column_values/u64_based/blockwise_linear.rs b/columnar/src/column_values/u64_based/blockwise_linear.rs index b60bf5bad..b8d03794f 100644 --- a/columnar/src/column_values/u64_based/blockwise_linear.rs +++ b/columnar/src/column_values/u64_based/blockwise_linear.rs @@ -11,7 +11,7 @@ use crate::column_values::u64_based::line::Line; use crate::column_values::u64_based::{ColumnCodec, ColumnCodecEstimator, ColumnStats}; use crate::column_values::{ColumnValues, VecColumn}; -const BLOCK_SIZE: u32 = 512u32; +pub(crate) const BLOCK_SIZE: u32 = 512u32; #[derive(Debug, Default)] struct Block { diff --git a/columnar/src/column_values/u64_based/mod.rs b/columnar/src/column_values/u64_based/mod.rs index aa2d9818b..42ef335ad 100644 --- a/columnar/src/column_values/u64_based/mod.rs +++ b/columnar/src/column_values/u64_based/mod.rs @@ -1,5 +1,5 @@ mod bitpacked; -mod blockwise_linear; +pub(crate) mod blockwise_linear; mod line; mod linear; mod stats_collector; diff --git a/columnar/src/columnar/merge/mod.rs b/columnar/src/columnar/merge/mod.rs index 4f7739f4d..c1a73a97a 100644 --- a/columnar/src/columnar/merge/mod.rs +++ b/columnar/src/columnar/merge/mod.rs @@ -79,6 +79,23 @@ pub fn merge_columnar( required_columns: &[(String, ColumnType)], merge_row_order: MergeRowOrder, output: &mut impl io::Write, +) -> io::Result<()> { + merge_columnar_with_tie_breakers( + columnar_readers, + required_columns, + &[], + merge_row_order, + output, + ) +} + +/// Merges columnars and regenerates the named tie-breaker columns in the resulting row order. +pub fn merge_columnar_with_tie_breakers( + columnar_readers: &[&ColumnarReader], + required_columns: &[(String, ColumnType)], + tie_breaker_columns: &[String], + merge_row_order: MergeRowOrder, + output: &mut impl io::Write, ) -> io::Result<()> { let mut serializer = ColumnarSerializer::new(output); let num_docs_per_columnar = columnar_readers @@ -89,6 +106,17 @@ pub fn merge_columnar( let columns_to_merge = group_columns_for_merge(columnar_readers, required_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 mut column_serializer = + serializer.start_serialize_column(column_name.as_bytes(), ColumnType::U64); + crate::column::serialize_generated_tie_breaker_column( + merge_row_order.num_rows(), + &mut column_serializer, + )?; + column_serializer.finalize()?; + continue; + } + let grouped_columns = grouped_columns.open(&merge_row_order)?; if grouped_columns.is_empty() { continue; diff --git a/columnar/src/columnar/mod.rs b/columnar/src/columnar/mod.rs index 0c23f9e3b..d98fa70c3 100644 --- a/columnar/src/columnar/mod.rs +++ b/columnar/src/columnar/mod.rs @@ -10,7 +10,7 @@ pub use format_version::{CURRENT_VERSION, Version}; pub(crate) use merge::ColumnTypeCategory; pub use merge::{ MergeRowOrder, ShuffleMergeOrder, StackMergeOrder, compute_merged_term_ord_mapping, - merge_columnar, + merge_columnar, merge_columnar_with_tie_breakers, }; pub use reader::ColumnarReader; pub use writer::ColumnarWriter; diff --git a/columnar/src/columnar/reader/mod.rs b/columnar/src/columnar/reader/mod.rs index 8592b3a22..30734ea79 100644 --- a/columnar/src/columnar/reader/mod.rs +++ b/columnar/src/columnar/reader/mod.rs @@ -236,6 +236,25 @@ mod tests { assert_eq!(columns[1].1.column_type(), ColumnType::U64); } + #[test] + fn test_generated_tie_breaker_column() { + let mut columnar_writer = ColumnarWriter::default(); + columnar_writer.record_tie_breaker_column("tie"); + let mut buffer = Vec::new(); + columnar_writer.serialize(1_025, None, &mut buffer).unwrap(); + + let columnar = ColumnarReader::open(buffer).unwrap(); + let handles = columnar.read_columns("tie").unwrap(); + let column = handles[0].open_u64_lenient().unwrap().unwrap(); + assert_eq!(column.index.get_cardinality(), crate::Cardinality::Full); + for block_start in [0, 512, 1_024] { + let block_end = (block_start + 512).min(1_025); + for doc in block_start + 1..block_end { + assert_eq!(column.first(doc), Some(column.first(doc - 1).unwrap() + 1)); + } + } + } + #[test] fn test_list_columns_strict_typing_prevents_coercion() { let mut columnar_writer = ColumnarWriter::default(); diff --git a/columnar/src/columnar/writer/mod.rs b/columnar/src/columnar/writer/mod.rs index 999ccd058..049a8a6f4 100644 --- a/columnar/src/columnar/writer/mod.rs +++ b/columnar/src/columnar/writer/mod.rs @@ -3,6 +3,7 @@ mod column_writers; mod serializer; mod value_index; +use std::collections::HashSet; use std::io; use std::net::Ipv6Addr; @@ -54,6 +55,7 @@ pub struct ColumnarWriter { ip_addr_field_hash_map: ArenaHashMap, bytes_field_hash_map: ArenaHashMap, str_field_hash_map: ArenaHashMap, + generated_tie_breaker_columns: HashSet>, arena: MemoryArena, // Dictionaries used to store dictionary-encoded values. dictionaries: Vec, @@ -217,6 +219,13 @@ impl ColumnarWriter { } } + /// Registers a full `u64` column whose values are generated when this writer is serialized. + pub fn record_tie_breaker_column(&mut self, column_name: &str) { + self.record_column_type(column_name, ColumnType::U64, false); + self.generated_tie_breaker_columns + .insert(column_name.as_bytes().to_vec()); + } + pub fn record_numerical + Copy>( &mut self, doc: RowId, @@ -436,24 +445,31 @@ impl ColumnarWriter { column_serializer.finalize()?; } ColumnType::F64 | ColumnType::I64 | ColumnType::U64 => { - let numerical_column_writer: NumericalColumnWriter = - self.numerical_field_hash_map.read(addr); - let cardinality = numerical_column_writer.cardinality(num_docs); let mut column_serializer = serializer.start_serialize_column(column_name, column_type); - let numerical_type = column_type.numerical_type().unwrap(); - serialize_numerical_column( - cardinality, - num_docs, - numerical_type, - numerical_column_writer.operation_iterator( - arena, - old_to_new_row_ids, - &mut symbol_byte_buffer, - ), - buffers, - &mut column_serializer, - )?; + if self.generated_tie_breaker_columns.contains(column_name) { + crate::column::serialize_generated_tie_breaker_column( + num_docs, + &mut column_serializer, + )?; + } else { + let numerical_column_writer: NumericalColumnWriter = + self.numerical_field_hash_map.read(addr); + let cardinality = numerical_column_writer.cardinality(num_docs); + let numerical_type = column_type.numerical_type().unwrap(); + serialize_numerical_column( + cardinality, + num_docs, + numerical_type, + numerical_column_writer.operation_iterator( + arena, + old_to_new_row_ids, + &mut symbol_byte_buffer, + ), + buffers, + &mut column_serializer, + )?; + } column_serializer.finalize()?; } ColumnType::DateTime => { diff --git a/columnar/src/lib.rs b/columnar/src/lib.rs index 1da8d9604..60e0e54ea 100644 --- a/columnar/src/lib.rs +++ b/columnar/src/lib.rs @@ -42,7 +42,7 @@ pub use column_values::{ pub use columnar::{ CURRENT_VERSION, ColumnType, ColumnarReader, ColumnarWriter, HasAssociatedColumnType, MergeRowOrder, ShuffleMergeOrder, StackMergeOrder, Version, compute_merged_term_ord_mapping, - merge_columnar, + merge_columnar, merge_columnar_with_tie_breakers, }; use sstable::VoidSSTable; pub use value::{NumericalType, NumericalValue}; diff --git a/src/fastfield/mod.rs b/src/fastfield/mod.rs index d56dc27a8..68bb0e758 100644 --- a/src/fastfield/mod.rs +++ b/src/fastfield/mod.rs @@ -148,6 +148,28 @@ mod tests { Ok(()) } + #[test] + fn test_generated_tie_breaker_fast_field() { + let mut schema_builder = Schema::builder(); + schema_builder.add_tie_breaker_field("tie"); + let schema = schema_builder.build(); + let mut writer = FastFieldsWriter::from_schema(&schema).unwrap(); + for _ in 0..1_025 { + writer.add_document(&TantivyDocument::default()).unwrap(); + } + let mut bytes = Vec::new(); + writer.serialize(&mut bytes, None).unwrap(); + + let readers = FastFieldReaders::open(bytes.into(), schema).unwrap(); + let values = readers.u64("tie").unwrap().first_or_default_col(0); + for block_start in [0, 512, 1_024] { + let block_end = (block_start + 512).min(1_025); + for doc in block_start + 1..block_end { + assert_eq!(values.get_val(doc), values.get_val(doc - 1) + 1); + } + } + } + #[test] fn test_intfastfield_large() { let path = Path::new("test"); diff --git a/src/fastfield/plugin.rs b/src/fastfield/plugin.rs index eb33b6816..37eabc4ae 100644 --- a/src/fastfield/plugin.rs +++ b/src/fastfield/plugin.rs @@ -18,7 +18,7 @@ use crate::index::{SegmentComponent, SegmentReader}; use crate::indexer::doc_id_mapping::{DocIdMapping, MappingType, SegmentDocIdMapping}; use crate::plugin::{PluginMergeContext, PluginWriter, PluginWriterContext, SegmentPlugin}; use crate::schema::document::Document; -use crate::schema::{value_type_to_column_type, Schema}; +use crate::schema::{value_type_to_column_type, FieldType, Schema}; use crate::space_usage::{ComponentSpaceUsage, FAST_FIELDS}; use crate::Segment; @@ -43,6 +43,7 @@ impl SegmentPlugin for FastFieldsPlugin { ctx.target_segment.index().directory().open_write(&path)?; let required_columns = extract_fast_field_required_columns(ctx.schema); + let tie_breaker_columns = extract_tie_breaker_columns(ctx.schema); let columnars: Vec<&ColumnarReader> = ctx .readers .iter() @@ -53,9 +54,10 @@ impl SegmentPlugin for FastFieldsPlugin { let doc_id_mapping = ctx.doc_id_mapping.clone(); let merge_row_order = convert_to_merge_order(&columnars[..], doc_id_mapping); - columnar::merge_columnar( + columnar::merge_columnar_with_tie_breakers( &columnars[..], &required_columns, + &tie_breaker_columns, merge_row_order, &mut fast_field_wrt, )?; @@ -173,6 +175,18 @@ fn convert_to_merge_order( } } +fn extract_tie_breaker_columns(schema: &Schema) -> Vec { + schema + .fields() + .filter_map(|(_, field_entry)| match field_entry.field_type() { + FieldType::U64(options) if options.is_tie_breaker() => { + Some(field_entry.name().to_string()) + } + _ => None, + }) + .collect() +} + fn extract_fast_field_required_columns(schema: &Schema) -> Vec<(String, ColumnType)> { schema .fields() diff --git a/src/fastfield/writer.rs b/src/fastfield/writer.rs index 9bca41357..d74b68ae4 100644 --- a/src/fastfield/writer.rs +++ b/src/fastfield/writer.rs @@ -52,6 +52,13 @@ impl FastFieldsWriter { if !field_entry.field_type().is_fast() { continue; } + if matches!( + field_entry.field_type(), + FieldType::U64(options) if options.is_tie_breaker() + ) { + columnar_writer.record_tie_breaker_column(field_entry.name()); + continue; + } fast_field_names[field_id.field_id() as usize] = Some(field_entry.name().to_string()); let value_type = field_entry.field_type().value_type(); if let FieldType::Date(date_options) = field_entry.field_type() { diff --git a/src/schema/numeric_options.rs b/src/schema/numeric_options.rs index db36b523e..39f6fffd8 100644 --- a/src/schema/numeric_options.rs +++ b/src/schema/numeric_options.rs @@ -16,6 +16,8 @@ pub struct NumericOptions { stored: bool, #[serde(skip_serializing_if = "is_false")] coerce: bool, + #[serde(skip_serializing_if = "is_false")] + tie_breaker: bool, } fn is_false(val: &bool) -> bool { @@ -37,6 +39,8 @@ struct NumericOptionsDeser { stored: bool, #[serde(default)] coerce: bool, + #[serde(default)] + tie_breaker: bool, } impl From for NumericOptions { @@ -47,6 +51,7 @@ impl From for NumericOptions { fast: deser.fast, stored: deser.stored, coerce: deser.coerce, + tie_breaker: deser.tie_breaker, } } } @@ -82,6 +87,16 @@ impl NumericOptions { self.coerce } + pub(crate) fn is_tie_breaker(&self) -> bool { + self.tie_breaker + } + + pub(crate) fn set_tie_breaker(mut self) -> Self { + self.fast = true; + self.tie_breaker = true; + self + } + /// Try to coerce values if they are not a number. Defaults to false. #[must_use] pub fn set_coerce(mut self) -> Self { @@ -145,6 +160,7 @@ impl From for NumericOptions { stored: false, fast: false, coerce: true, + tie_breaker: false, } } } @@ -157,6 +173,7 @@ impl From for NumericOptions { stored: false, fast: true, coerce: false, + tie_breaker: false, } } } @@ -169,6 +186,7 @@ impl From for NumericOptions { stored: true, fast: false, coerce: false, + tie_breaker: false, } } } @@ -181,6 +199,7 @@ impl From for NumericOptions { stored: false, fast: false, coerce: false, + tie_breaker: false, } } } @@ -196,6 +215,7 @@ impl> BitOr for NumericOptions { stored: self.stored | other.stored, fast: self.fast | other.fast, coerce: self.coerce | other.coerce, + tie_breaker: self.tie_breaker | other.tie_breaker, } } } @@ -230,6 +250,7 @@ mod tests { fast: false, stored: false, coerce: false, + tie_breaker: false, } ); } @@ -249,6 +270,7 @@ mod tests { fast: false, stored: false, coerce: false, + tie_breaker: false, } ); } @@ -269,6 +291,7 @@ mod tests { fast: false, stored: false, coerce: false, + tie_breaker: false, } ); } @@ -290,6 +313,7 @@ mod tests { fast: false, stored: false, coerce: false, + tie_breaker: false, } ); } @@ -312,6 +336,7 @@ mod tests { fast: false, stored: false, coerce: true, + tie_breaker: false, } ); } diff --git a/src/schema/schema.rs b/src/schema/schema.rs index 79414473e..833174847 100644 --- a/src/schema/schema.rs +++ b/src/schema/schema.rs @@ -57,6 +57,18 @@ impl SchemaBuilder { self.add_field(field_entry) } + /// Adds a generated tie-breaker fast field. + /// + /// The field is exposed as a `u64` fast field, but its generated values fit in a `u32`. + /// Values supplied by documents for this field are ignored. + /// + /// # Panics + /// + /// Panics when field already exists. + pub fn add_tie_breaker_field(&mut self, field_name_str: &str) -> Field { + self.add_u64_field(field_name_str, NumericOptions::default().set_tie_breaker()) + } + /// Adds a new i64 field. /// Returns the associated field handle ///