Compare commits

...

3 Commits

Author SHA1 Message Date
Paul Masurel
b91c37ab2b bumping columnar version 2024-06-11 16:38:10 +09:00
Paul Masurel
ec7379477c low hanging fruit in optimization 2024-06-11 16:25:20 +09:00
Paul Masurel
31b78b59dd DONOTMERGE: Hack placing a optional index in front of multivalued indexes.
The point is to avoid the large overhead added to computing the start
offset index when a column is for the most part sparse.
2024-06-11 16:22:49 +09:00
18 changed files with 802 additions and 560 deletions

View File

@@ -150,61 +150,62 @@ mod tests {
); );
} }
#[test] // #[test]
fn test_merge_index_multivalued_sorted() { // fn test_merge_index_multivalued_sorted() {
let column_indexes: Vec<ColumnIndex> = vec![MultiValueIndex::for_test(&[0, 2, 5]).into()]; // let column_indexes: Vec<ColumnIndex> = vec![MultiValueIndex::for_test(&[0, 2,
let merge_row_order: MergeRowOrder = ShuffleMergeOrder::for_test( // 5]).into()]; let merge_row_order: MergeRowOrder = ShuffleMergeOrder::for_test(
&[2], // &[2],
vec![ // vec![
RowAddr { // RowAddr {
segment_ord: 0u32, // segment_ord: 0u32,
row_id: 1u32, // row_id: 1u32,
}, // },
RowAddr { // RowAddr {
segment_ord: 0u32, // segment_ord: 0u32,
row_id: 0u32, // row_id: 0u32,
}, // },
], // ],
) // )
.into(); // .into();
let merged_column_index = merge_column_index(&column_indexes[..], &merge_row_order); // let merged_column_index = merge_column_index(&column_indexes[..], &merge_row_order);
let SerializableColumnIndex::Multivalued(start_index_iterable) = merged_column_index else { // let SerializableColumnIndex::Multivalued(serializable_multivalue_index) =
panic!("Excpected a multivalued index") // merged_column_index else { panic!("Excpected a multivalued index")
}; // };
let start_indexes: Vec<RowId> = start_index_iterable.boxed_iter().collect(); // serializable_multivalue_index.doc_ids_with_values_opt.
assert_eq!(&start_indexes, &[0, 3, 5]); // let start_indexes: Vec<RowId> = start_index_iterable.boxed_iter().collect();
} // assert_eq!(&start_indexes, &[0, 3, 5]);
// }
#[test] // #[test]
fn test_merge_index_multivalued_sorted_several_segment() { // fn test_merge_index_multivalued_sorted_several_segment() {
let column_indexes: Vec<ColumnIndex> = vec![ // let column_indexes: Vec<ColumnIndex> = vec![
MultiValueIndex::for_test(&[0, 2, 5]).into(), // MultiValueIndex::for_test(&[0, 2, 5]).into(),
ColumnIndex::Empty { num_docs: 0 }, // ColumnIndex::Empty { num_docs: 0 },
MultiValueIndex::for_test(&[0, 1, 4]).into(), // MultiValueIndex::for_test(&[0, 1, 4]).into(),
]; // ];
let merge_row_order: MergeRowOrder = ShuffleMergeOrder::for_test( // let merge_row_order: MergeRowOrder = ShuffleMergeOrder::for_test(
&[2, 0, 2], // &[2, 0, 2],
vec![ // vec![
RowAddr { // RowAddr {
segment_ord: 2u32, // segment_ord: 2u32,
row_id: 1u32, // row_id: 1u32,
}, // },
RowAddr { // RowAddr {
segment_ord: 0u32, // segment_ord: 0u32,
row_id: 0u32, // row_id: 0u32,
}, // },
RowAddr { // RowAddr {
segment_ord: 2u32, // segment_ord: 2u32,
row_id: 0u32, // row_id: 0u32,
}, // },
], // ],
) // )
.into(); // .into();
let merged_column_index = merge_column_index(&column_indexes[..], &merge_row_order); // let merged_column_index = merge_column_index(&column_indexes[..], &merge_row_order);
let SerializableColumnIndex::Multivalued(start_index_iterable) = merged_column_index else { // let SerializableColumnIndex::Multivalued(serializable_multivalue_index) =
panic!("Excpected a multivalued index") // merged_column_index else { panic!("Excpected a multivalued index")
}; // };
let start_indexes: Vec<RowId> = start_index_iterable.boxed_iter().collect(); // let start_indexes: Vec<RowId> = start_index_iterable.boxed_iter().collect();
assert_eq!(&start_indexes, &[0, 3, 5, 6]); // assert_eq!(&start_indexes, &[0, 3, 5, 6]);
} // }
} }

View File

@@ -9,22 +9,23 @@ pub fn merge_column_index_shuffled<'a>(
cardinality_after_merge: Cardinality, cardinality_after_merge: Cardinality,
shuffle_merge_order: &'a ShuffleMergeOrder, shuffle_merge_order: &'a ShuffleMergeOrder,
) -> SerializableColumnIndex<'a> { ) -> SerializableColumnIndex<'a> {
match cardinality_after_merge { todo!();
Cardinality::Full => SerializableColumnIndex::Full, // match cardinality_after_merge {
Cardinality::Optional => { // Cardinality::Full => SerializableColumnIndex::Full,
let non_null_row_ids = // Cardinality::Optional => {
merge_column_index_shuffled_optional(column_indexes, shuffle_merge_order); // let non_null_row_ids =
SerializableColumnIndex::Optional { // merge_column_index_shuffled_optional(column_indexes, shuffle_merge_order);
non_null_row_ids, // SerializableColumnIndex::Optional {
num_rows: shuffle_merge_order.num_rows(), // non_null_row_ids,
} // num_rows: shuffle_merge_order.num_rows(),
} // }
Cardinality::Multivalued => { // }
let multivalue_start_index = // Cardinality::Multivalued => {
merge_column_index_shuffled_multivalued(column_indexes, shuffle_merge_order); // let multivalue_start_index =
SerializableColumnIndex::Multivalued(multivalue_start_index) // merge_column_index_shuffled_multivalued(column_indexes, shuffle_merge_order);
} // SerializableColumnIndex::Multivalued(multivalue_start_index)
} // }
// }
} }
/// Merge several column indexes into one, ordering rows according to the merge_order passed as /// Merge several column indexes into one, ordering rows according to the merge_order passed as
@@ -137,35 +138,35 @@ mod tests {
assert!(integrate_num_vals([3, 0, 10, 20].into_iter()).eq([0, 3, 3, 13, 33].into_iter())); assert!(integrate_num_vals([3, 0, 10, 20].into_iter()).eq([0, 3, 3, 13, 33].into_iter()));
} }
#[test] // #[test]
fn test_merge_column_index_optional_shuffle() { // fn test_merge_column_index_optional_shuffle() {
let optional_index: ColumnIndex = OptionalIndex::for_test(2, &[0]).into(); // let optional_index: ColumnIndex = OptionalIndex::for_test(2, &[0]).into();
let column_indexes = [optional_index, ColumnIndex::Full]; // let column_indexes = [optional_index, ColumnIndex::Full];
let row_addrs = vec![ // let row_addrs = vec![
RowAddr { // RowAddr {
segment_ord: 0u32, // segment_ord: 0u32,
row_id: 1u32, // row_id: 1u32,
}, // },
RowAddr { // RowAddr {
segment_ord: 1u32, // segment_ord: 1u32,
row_id: 0u32, // row_id: 0u32,
}, // },
]; // ];
let shuffle_merge_order = ShuffleMergeOrder::for_test(&[2, 1], row_addrs); // let shuffle_merge_order = ShuffleMergeOrder::for_test(&[2, 1], row_addrs);
let serializable_index = merge_column_index_shuffled( // let serializable_index = merge_column_index_shuffled(
&column_indexes[..], // &column_indexes[..],
Cardinality::Optional, // Cardinality::Optional,
&shuffle_merge_order, // &shuffle_merge_order,
); // );
let SerializableColumnIndex::Optional { // let SerializableColumnIndex::Optional {
non_null_row_ids, // non_null_row_ids,
num_rows, // num_rows,
} = serializable_index // } = serializable_index
else { // else {
panic!() // panic!()
}; // };
assert_eq!(num_rows, 2); // assert_eq!(num_rows, 2);
let non_null_rows: Vec<RowId> = non_null_row_ids.boxed_iter().collect(); // let non_null_rows: Vec<RowId> = non_null_row_ids.boxed_iter().collect();
assert_eq!(&non_null_rows, &[1]); // assert_eq!(&non_null_rows, &[1]);
} // }
} }

View File

@@ -1,6 +1,8 @@
use std::iter; use std::ops::Range;
use crate::column_index::{SerializableColumnIndex, Set}; use crate::column_index::multivalued_index::SerializableMultivalueIndex;
use crate::column_index::serialize::SerializableOptionalIndex;
use crate::column_index::SerializableColumnIndex;
use crate::iterable::Iterable; use crate::iterable::Iterable;
use crate::{Cardinality, ColumnIndex, RowId, StackMergeOrder}; use crate::{Cardinality, ColumnIndex, RowId, StackMergeOrder};
@@ -15,23 +17,140 @@ pub fn merge_column_index_stacked<'a>(
) -> SerializableColumnIndex<'a> { ) -> SerializableColumnIndex<'a> {
match cardinality_after_merge { match cardinality_after_merge {
Cardinality::Full => SerializableColumnIndex::Full, Cardinality::Full => SerializableColumnIndex::Full,
Cardinality::Optional => SerializableColumnIndex::Optional { Cardinality::Optional => SerializableColumnIndex::Optional(SerializableOptionalIndex {
non_null_row_ids: Box::new(StackedOptionalIndex { non_null_row_ids: Box::new(StackedOptionalIndex {
columns, columns,
stack_merge_order, stack_merge_order,
}), }),
num_rows: stack_merge_order.num_rows(), num_rows: stack_merge_order.num_rows(),
}, }),
Cardinality::Multivalued => { Cardinality::Multivalued => {
let stacked_multivalued_index = StackedMultivaluedIndex { let serializable_multivalue_index =
columns, make_serializable_multivalued_index(columns, stack_merge_order);
stack_merge_order, SerializableColumnIndex::Multivalued(serializable_multivalue_index)
};
SerializableColumnIndex::Multivalued(Box::new(stacked_multivalued_index))
} }
} }
} }
struct StackedDocIdsWithValues<'a> {
column_indexes: &'a [ColumnIndex],
stack_merge_order: &'a StackMergeOrder,
}
impl Iterable<u32> for StackedDocIdsWithValues<'_> {
fn boxed_iter(&self) -> Box<dyn Iterator<Item = u32> + '_> {
Box::new((0..self.column_indexes.len()).flat_map(|i| {
let column_index = &self.column_indexes[i];
let doc_range = self.stack_merge_order.columnar_range(i);
get_doc_ids_with_values(column_index, doc_range)
}))
}
}
fn get_doc_ids_with_values<'a>(
column_index: &'a ColumnIndex,
doc_range: Range<u32>,
) -> Box<dyn Iterator<Item = u32> + 'a> {
match column_index {
ColumnIndex::Empty { .. } => Box::new(0..0),
ColumnIndex::Full => Box::new(doc_range),
ColumnIndex::Optional(optional_index) => Box::new(
optional_index
.iter_rows()
.map(move |row| row + doc_range.start),
),
ColumnIndex::Multivalued(multivalued_index) => Box::new(
multivalued_index
.optional_index
.iter_rows()
.map(move |row| row + doc_range.start),
),
}
}
fn stack_doc_ids_with_values<'a>(
column_indexes: &'a [ColumnIndex],
stack_merge_order: &'a StackMergeOrder,
) -> SerializableOptionalIndex<'a> {
let num_rows = stack_merge_order.num_rows();
SerializableOptionalIndex {
non_null_row_ids: Box::new(StackedDocIdsWithValues {
column_indexes,
stack_merge_order,
}),
num_rows,
}
}
struct StackedStartOffsets<'a> {
column_indexes: &'a [ColumnIndex],
stack_merge_order: &'a StackMergeOrder,
}
fn get_num_values_iterator<'a>(
column_index: &'a ColumnIndex,
num_docs: u32,
) -> Box<dyn Iterator<Item = u32> + 'a> {
match column_index {
ColumnIndex::Empty { .. } => Box::new(std::iter::empty()),
ColumnIndex::Full => Box::new(std::iter::repeat(1u32).take(num_docs as usize)),
ColumnIndex::Optional(optional_index) => {
Box::new(std::iter::repeat(1u32).take(optional_index.num_non_nulls() as usize))
}
ColumnIndex::Multivalued(multivalued_index) => {
let vals: Vec<u32> = multivalued_index.start_index_column.iter().collect();
Box::new(
multivalued_index
.start_index_column
.iter()
.scan(0u32, |previous_start_offset, current_start_offset| {
let num_vals = current_start_offset - *previous_start_offset;
*previous_start_offset = current_start_offset;
Some(num_vals)
})
.skip(1),
)
}
}
}
impl<'a> Iterable for StackedStartOffsets<'a> {
fn boxed_iter(&self) -> Box<dyn Iterator<Item = u64> + '_> {
let num_values_it = (0..self.column_indexes.len()).flat_map(|columnar_id| {
let num_docs = self.stack_merge_order.columnar_range(columnar_id).len() as u32;
let column_index = &self.column_indexes[columnar_id];
get_num_values_iterator(column_index, num_docs)
});
Box::new(std::iter::once(0u64).chain(num_values_it.into_iter().scan(
0u64,
|cumulated, el| {
*cumulated += el as u64;
Some(*cumulated)
},
)))
}
}
fn stack_start_offsets<'a>(
column_indexes: &'a [ColumnIndex],
stack_merge_order: &'a StackMergeOrder,
) -> Box<dyn Iterable + 'a> {
Box::new(StackedStartOffsets {
column_indexes,
stack_merge_order,
})
}
fn make_serializable_multivalued_index<'a>(
columns: &'a [ColumnIndex],
stack_merge_order: &'a StackMergeOrder,
) -> SerializableMultivalueIndex<'a> {
SerializableMultivalueIndex {
doc_ids_with_values: stack_doc_ids_with_values(columns, stack_merge_order),
start_offsets: stack_start_offsets(columns, stack_merge_order),
}
}
struct StackedOptionalIndex<'a> { struct StackedOptionalIndex<'a> {
columns: &'a [ColumnIndex], columns: &'a [ColumnIndex],
stack_merge_order: &'a StackMergeOrder, stack_merge_order: &'a StackMergeOrder,
@@ -62,87 +181,3 @@ impl<'a> Iterable<RowId> for StackedOptionalIndex<'a> {
) )
} }
} }
#[derive(Clone, Copy)]
struct StackedMultivaluedIndex<'a> {
columns: &'a [ColumnIndex],
stack_merge_order: &'a StackMergeOrder,
}
fn convert_column_opt_to_multivalued_index<'a>(
column_index_opt: &'a ColumnIndex,
num_rows: RowId,
) -> Box<dyn Iterator<Item = RowId> + 'a> {
match column_index_opt {
ColumnIndex::Empty { .. } => Box::new(iter::repeat(0u32).take(num_rows as usize + 1)),
ColumnIndex::Full => Box::new(0..num_rows + 1),
ColumnIndex::Optional(optional_index) => {
Box::new(
(0..num_rows)
// TODO optimize
.map(|row_id| optional_index.rank(row_id))
.chain(std::iter::once(optional_index.num_non_nulls())),
)
}
ColumnIndex::Multivalued(multivalued_index) => multivalued_index.start_index_column.iter(),
}
}
impl<'a> Iterable<RowId> for StackedMultivaluedIndex<'a> {
fn boxed_iter(&self) -> Box<dyn Iterator<Item = RowId> + '_> {
let multivalued_indexes =
self.columns
.iter()
.enumerate()
.map(|(columnar_id, column_opt)| {
let num_rows =
self.stack_merge_order.columnar_range(columnar_id).len() as RowId;
convert_column_opt_to_multivalued_index(column_opt, num_rows)
});
stack_multivalued_indexes(multivalued_indexes)
}
}
// Refactor me
fn stack_multivalued_indexes<'a>(
mut multivalued_indexes: impl Iterator<Item = Box<dyn Iterator<Item = RowId> + 'a>> + 'a,
) -> Box<dyn Iterator<Item = RowId> + 'a> {
let mut offset = 0;
let mut last_row_id = 0;
let mut current_it = multivalued_indexes.next();
Box::new(std::iter::from_fn(move || loop {
if let Some(row_id) = current_it.as_mut()?.next() {
last_row_id = offset + row_id;
return Some(last_row_id);
}
offset = last_row_id;
loop {
current_it = multivalued_indexes.next();
if current_it.as_mut()?.next().is_some() {
break;
}
}
}))
}
#[cfg(test)]
mod tests {
use crate::RowId;
fn it<'a>(row_ids: &'a [RowId]) -> Box<dyn Iterator<Item = RowId> + 'a> {
Box::new(row_ids.iter().copied())
}
#[test]
fn test_stack() {
let columns = [
it(&[0u32, 0u32]),
it(&[0u32, 1u32, 1u32, 4u32]),
it(&[0u32, 3u32, 5u32]),
it(&[0u32, 4u32]),
]
.into_iter();
let start_offsets: Vec<RowId> = super::stack_multivalued_indexes(columns).collect();
assert_eq!(start_offsets, &[0, 0, 1, 1, 4, 7, 9, 13]);
}
}

View File

@@ -11,8 +11,11 @@ mod serialize;
use std::ops::Range; use std::ops::Range;
pub use merge::merge_column_index; pub use merge::merge_column_index;
pub(crate) use multivalued_index::SerializableMultivalueIndex;
pub use optional_index::{OptionalIndex, Set}; pub use optional_index::{OptionalIndex, Set};
pub use serialize::{open_column_index, serialize_column_index, SerializableColumnIndex}; pub use serialize::{
open_column_index, serialize_column_index, SerializableColumnIndex, SerializableOptionalIndex,
};
use crate::column_index::multivalued_index::MultiValueIndex; use crate::column_index::multivalued_index::MultiValueIndex;
use crate::{Cardinality, DocId, RowId}; use crate::{Cardinality, DocId, RowId};

View File

@@ -3,35 +3,64 @@ use std::io::Write;
use std::ops::Range; use std::ops::Range;
use std::sync::Arc; use std::sync::Arc;
use common::OwnedBytes; use common::{CountingWriter, OwnedBytes};
use super::optional_index::{open_optional_index, serialize_optional_index};
use super::{OptionalIndex, SerializableOptionalIndex, Set};
use crate::column_values::{ use crate::column_values::{
load_u64_based_column_values, serialize_u64_based_column_values, CodecType, ColumnValues, load_u64_based_column_values, serialize_u64_based_column_values, CodecType, ColumnValues,
}; };
use crate::iterable::Iterable; use crate::iterable::Iterable;
use crate::{DocId, RowId}; use crate::{DocId, RowId};
pub fn serialize_multivalued_index( pub struct SerializableMultivalueIndex<'a> {
multivalued_index: &dyn Iterable<RowId>, pub doc_ids_with_values: SerializableOptionalIndex<'a>,
pub start_offsets: Box<dyn Iterable<u64> + 'a>,
}
pub fn serialize_multivalued_index<'a>(
multivalued_index: &SerializableMultivalueIndex<'a>,
output: &mut impl Write, output: &mut impl Write,
) -> io::Result<()> { ) -> io::Result<()> {
let SerializableMultivalueIndex {
doc_ids_with_values,
start_offsets,
} = multivalued_index;
let mut count_writer = CountingWriter::wrap(output);
let SerializableOptionalIndex {
non_null_row_ids,
num_rows,
} = doc_ids_with_values;
serialize_optional_index(&**non_null_row_ids, *num_rows, &mut count_writer)?;
let optional_len = count_writer.written_bytes() as u32;
let output = count_writer.finish();
serialize_u64_based_column_values( serialize_u64_based_column_values(
multivalued_index, &**start_offsets,
&[CodecType::Bitpacked, CodecType::Linear], &[CodecType::Bitpacked, CodecType::Linear],
output, output,
)?; )?;
output.write_all(&optional_len.to_le_bytes())?;
Ok(()) Ok(())
} }
pub fn open_multivalued_index(bytes: OwnedBytes) -> io::Result<MultiValueIndex> { pub fn open_multivalued_index(bytes: OwnedBytes) -> io::Result<MultiValueIndex> {
let start_index_column: Arc<dyn ColumnValues<RowId>> = load_u64_based_column_values(bytes)?; let (body_bytes, optional_index_len) = bytes.rsplit(4);
Ok(MultiValueIndex { start_index_column }) let optional_index_len = u32::from_le_bytes(optional_index_len.as_slice().try_into().unwrap());
let (optional_index_bytes, start_index_bytes) = body_bytes.split(optional_index_len as usize);
let optional_index = open_optional_index(optional_index_bytes)?;
let start_index_column: Arc<dyn ColumnValues<RowId>> =
load_u64_based_column_values(start_index_bytes)?;
Ok(MultiValueIndex {
optional_index,
start_index_column,
})
} }
#[derive(Clone)] #[derive(Clone)]
/// Index to resolve value range for given doc_id. /// Index to resolve value range for given doc_id.
/// Starts at 0. /// Starts at 0.
pub struct MultiValueIndex { pub struct MultiValueIndex {
pub optional_index: OptionalIndex,
pub start_index_column: Arc<dyn crate::ColumnValues<RowId>>, pub start_index_column: Arc<dyn crate::ColumnValues<RowId>>,
} }
@@ -43,16 +72,27 @@ impl std::fmt::Debug for MultiValueIndex {
} }
} }
impl From<Arc<dyn ColumnValues<RowId>>> for MultiValueIndex {
fn from(start_index_column: Arc<dyn ColumnValues<RowId>>) -> Self {
MultiValueIndex { start_index_column }
}
}
impl MultiValueIndex { impl MultiValueIndex {
pub fn for_test(start_offsets: &[RowId]) -> MultiValueIndex { pub fn for_test(start_offsets: &[RowId]) -> MultiValueIndex {
assert!(start_offsets.len() > 0);
assert_eq!(start_offsets[0], 0);
let mut doc_with_values = Vec::new();
let mut compact_start_offsets: Vec<u64> = vec![0];
for doc in 0..start_offsets.len() - 1 {
if start_offsets[doc] < start_offsets[doc + 1] {
doc_with_values.push(doc as RowId);
compact_start_offsets.push(start_offsets[doc + 1] as u64);
}
}
let serializable_multivalued_index = SerializableMultivalueIndex {
doc_ids_with_values: SerializableOptionalIndex {
non_null_row_ids: Box::new(&doc_with_values[..]),
num_rows: start_offsets.len() as u32 - 1,
},
start_offsets: Box::new(&compact_start_offsets[..]),
};
let mut buffer = Vec::new(); let mut buffer = Vec::new();
serialize_multivalued_index(&start_offsets, &mut buffer).unwrap(); serialize_multivalued_index(&serializable_multivalued_index, &mut buffer).unwrap();
let bytes = OwnedBytes::new(buffer); let bytes = OwnedBytes::new(buffer);
open_multivalued_index(bytes).unwrap() open_multivalued_index(bytes).unwrap()
} }
@@ -61,15 +101,19 @@ impl MultiValueIndex {
/// the given document are `start..end`. /// the given document are `start..end`.
#[inline] #[inline]
pub(crate) fn range(&self, doc_id: DocId) -> Range<RowId> { pub(crate) fn range(&self, doc_id: DocId) -> Range<RowId> {
let start = self.start_index_column.get_val(doc_id); let Some(rank) = self.optional_index.rank_if_exists(doc_id) else {
let end = self.start_index_column.get_val(doc_id + 1); return 0..0;
};
let start = self.start_index_column.get_val(rank);
let end = self.start_index_column.get_val(rank + 1);
start..end start..end
} }
/// Returns the number of documents in the index. /// Returns the number of documents in the index.
#[inline] #[inline]
pub fn num_docs(&self) -> u32 { pub fn num_docs(&self) -> u32 {
self.start_index_column.num_vals() - 1 self.optional_index.num_docs()
// self.start_index_column.num_vals() - 1
} }
/// Converts a list of ranks (row ids of values) in a 1:n index to the corresponding list of /// Converts a list of ranks (row ids of values) in a 1:n index to the corresponding list of
@@ -108,6 +152,10 @@ impl MultiValueIndex {
} }
} }
ranks.truncate(write_doc_pos); ranks.truncate(write_doc_pos);
for rank in ranks.iter_mut() {
*rank = self.optional_index.select(*rank);
}
} }
} }
@@ -134,6 +182,7 @@ mod tests {
let positions = &[10u32, 11, 15, 20, 21, 22]; let positions = &[10u32, 11, 15, 20, 21, 22];
assert_eq!(index_to_pos_helper(&index, 0..5, positions), vec![1, 3, 4]); assert_eq!(index_to_pos_helper(&index, 0..5, positions), vec![1, 3, 4]);
assert_eq!(index_to_pos_helper(&index, 1..5, positions), vec![1, 3, 4]); assert_eq!(index_to_pos_helper(&index, 1..5, positions), vec![1, 3, 4]);
assert_eq!(index_to_pos_helper(&index, 0..5, &[9]), vec![0]); assert_eq!(index_to_pos_helper(&index, 0..5, &[9]), vec![0]);
assert_eq!(index_to_pos_helper(&index, 1..5, &[10]), vec![1]); assert_eq!(index_to_pos_helper(&index, 1..5, &[10]), vec![1]);
assert_eq!(index_to_pos_helper(&index, 1..5, &[11]), vec![1]); assert_eq!(index_to_pos_helper(&index, 1..5, &[11]), vec![1]);

View File

@@ -86,8 +86,14 @@ pub struct OptionalIndex {
block_metas: Arc<[BlockMeta]>, block_metas: Arc<[BlockMeta]>,
} }
impl<'a> Iterable<u32> for &'a OptionalIndex {
fn boxed_iter(&self) -> Box<dyn Iterator<Item = u32> + '_> {
Box::new(self.iter_rows())
}
}
impl std::fmt::Debug for OptionalIndex { impl std::fmt::Debug for OptionalIndex {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
f.debug_struct("OptionalIndex") f.debug_struct("OptionalIndex")
.field("num_rows", &self.num_rows) .field("num_rows", &self.num_rows)
.field("num_non_null_rows", &self.num_non_null_rows) .field("num_non_null_rows", &self.num_non_null_rows)
@@ -250,6 +256,10 @@ impl Set<RowId> for OptionalIndex {
} }
impl OptionalIndex { impl OptionalIndex {
pub fn new_empty(num_rows: RowId) -> OptionalIndex {
Self::for_test(num_rows, &[])
}
pub fn for_test(num_rows: RowId, row_ids: &[RowId]) -> OptionalIndex { pub fn for_test(num_rows: RowId, row_ids: &[RowId]) -> OptionalIndex {
assert!(row_ids assert!(row_ids
.last() .last()

View File

@@ -3,28 +3,41 @@ use std::io::Write;
use common::{CountingWriter, OwnedBytes}; use common::{CountingWriter, OwnedBytes};
use super::multivalued_index::SerializableMultivalueIndex;
use super::OptionalIndex;
use crate::column_index::multivalued_index::serialize_multivalued_index; use crate::column_index::multivalued_index::serialize_multivalued_index;
use crate::column_index::optional_index::serialize_optional_index; use crate::column_index::optional_index::serialize_optional_index;
use crate::column_index::ColumnIndex; use crate::column_index::ColumnIndex;
use crate::iterable::Iterable; use crate::iterable::Iterable;
use crate::{Cardinality, RowId}; use crate::{Cardinality, RowId};
pub struct SerializableOptionalIndex<'a> {
pub non_null_row_ids: Box<dyn Iterable<RowId> + 'a>,
pub num_rows: RowId,
}
impl<'a> From<&'a OptionalIndex> for SerializableOptionalIndex<'a> {
fn from(optional_index: &'a OptionalIndex) -> Self {
SerializableOptionalIndex {
non_null_row_ids: Box::new(optional_index),
num_rows: optional_index.num_docs(),
}
}
}
pub enum SerializableColumnIndex<'a> { pub enum SerializableColumnIndex<'a> {
Full, Full,
Optional { Optional(SerializableOptionalIndex<'a>),
non_null_row_ids: Box<dyn Iterable<RowId> + 'a>,
num_rows: RowId,
},
// TODO remove the Arc<dyn> apart from serialization this is not // TODO remove the Arc<dyn> apart from serialization this is not
// dynamic at all. // dynamic at all.
Multivalued(Box<dyn Iterable<RowId> + 'a>), Multivalued(SerializableMultivalueIndex<'a>),
} }
impl<'a> SerializableColumnIndex<'a> { impl<'a> SerializableColumnIndex<'a> {
pub fn get_cardinality(&self) -> Cardinality { pub fn get_cardinality(&self) -> Cardinality {
match self { match self {
SerializableColumnIndex::Full => Cardinality::Full, SerializableColumnIndex::Full => Cardinality::Full,
SerializableColumnIndex::Optional { .. } => Cardinality::Optional, SerializableColumnIndex::Optional(_) => Cardinality::Optional,
SerializableColumnIndex::Multivalued(_) => Cardinality::Multivalued, SerializableColumnIndex::Multivalued(_) => Cardinality::Multivalued,
} }
} }
@@ -40,12 +53,12 @@ pub fn serialize_column_index(
output.write_all(&[cardinality])?; output.write_all(&[cardinality])?;
match column_index { match column_index {
SerializableColumnIndex::Full => {} SerializableColumnIndex::Full => {}
SerializableColumnIndex::Optional { SerializableColumnIndex::Optional(SerializableOptionalIndex {
non_null_row_ids, non_null_row_ids,
num_rows, num_rows,
} => serialize_optional_index(non_null_row_ids.as_ref(), num_rows, &mut output)?, }) => serialize_optional_index(non_null_row_ids.as_ref(), num_rows, &mut output)?,
SerializableColumnIndex::Multivalued(multivalued_index) => { SerializableColumnIndex::Multivalued(multivalued_index) => {
serialize_multivalued_index(&*multivalued_index, &mut output)? serialize_multivalued_index(&multivalued_index, &mut output)?
} }
} }
let column_index_num_bytes = output.written_bytes() as u32; let column_index_num_bytes = output.written_bytes() as u32;

View File

@@ -8,7 +8,7 @@ const MAGIC_BYTES: [u8; 4] = [2, 113, 119, 66];
pub fn footer() -> [u8; VERSION_FOOTER_NUM_BYTES] { pub fn footer() -> [u8; VERSION_FOOTER_NUM_BYTES] {
let mut footer_bytes = [0u8; VERSION_FOOTER_NUM_BYTES]; let mut footer_bytes = [0u8; VERSION_FOOTER_NUM_BYTES];
footer_bytes[0..4].copy_from_slice(&Version::V1.to_bytes()); footer_bytes[0..4].copy_from_slice(&Version::V2.to_bytes());
footer_bytes[4..8].copy_from_slice(&MAGIC_BYTES[..]); footer_bytes[4..8].copy_from_slice(&MAGIC_BYTES[..]);
footer_bytes footer_bytes
} }
@@ -24,6 +24,7 @@ pub fn parse_footer(footer_bytes: [u8; VERSION_FOOTER_NUM_BYTES]) -> Result<Vers
#[repr(u32)] #[repr(u32)]
pub enum Version { pub enum Version {
V1 = 1u32, V1 = 1u32,
V2 = 2u32,
} }
impl Version { impl Version {
@@ -34,7 +35,7 @@ impl Version {
fn try_from_bytes(bytes: [u8; 4]) -> Result<Version, InvalidData> { fn try_from_bytes(bytes: [u8; 4]) -> Result<Version, InvalidData> {
let code = u32::from_le_bytes(bytes); let code = u32::from_le_bytes(bytes);
match code { match code {
1u32 => Ok(Version::V1), 2u32 => Ok(Version::V2),
_ => Err(InvalidData), _ => Err(InvalidData),
} }
} }
@@ -49,7 +50,7 @@ mod tests {
#[test] #[test]
fn test_footer_dserialization() { fn test_footer_dserialization() {
let parsed_version: Version = parse_footer(footer()).unwrap(); let parsed_version: Version = parse_footer(footer()).unwrap();
assert_eq!(Version::V1, parsed_version); assert_eq!(Version::V2, parsed_version);
} }
#[test] #[test]
@@ -63,7 +64,7 @@ mod tests {
for &i in &version_to_tests { for &i in &version_to_tests {
let version_res = Version::try_from_bytes(i.to_le_bytes()); let version_res = Version::try_from_bytes(i.to_le_bytes());
if let Ok(version) = version_res { if let Ok(version) = version_res {
assert_eq!(version, Version::V1); assert_eq!(version, Version::V2);
assert_eq!(version.to_bytes(), i.to_le_bytes()); assert_eq!(version.to_bytes(), i.to_le_bytes());
valid_versions.insert(i); valid_versions.insert(i);
} }

View File

@@ -12,7 +12,7 @@ use common::CountingWriter;
pub(crate) use serializer::ColumnarSerializer; pub(crate) use serializer::ColumnarSerializer;
use stacker::{Addr, ArenaHashMap, MemoryArena}; use stacker::{Addr, ArenaHashMap, MemoryArena};
use crate::column_index::SerializableColumnIndex; use crate::column_index::{SerializableColumnIndex, SerializableOptionalIndex};
use crate::column_values::{MonotonicallyMappableToU128, MonotonicallyMappableToU64}; use crate::column_values::{MonotonicallyMappableToU128, MonotonicallyMappableToU64};
use crate::columnar::column_type::ColumnType; use crate::columnar::column_type::ColumnType;
use crate::columnar::writer::column_writers::{ use crate::columnar::writer::column_writers::{
@@ -20,6 +20,7 @@ use crate::columnar::writer::column_writers::{
}; };
use crate::columnar::writer::value_index::{IndexBuilder, PreallocatedIndexBuilders}; use crate::columnar::writer::value_index::{IndexBuilder, PreallocatedIndexBuilders};
use crate::dictionary::{DictionaryBuilder, TermIdMapping, UnorderedId}; use crate::dictionary::{DictionaryBuilder, TermIdMapping, UnorderedId};
use crate::iterable::Iterable;
use crate::value::{Coerce, NumericalType, NumericalValue}; use crate::value::{Coerce, NumericalType, NumericalValue};
use crate::{Cardinality, RowId}; use crate::{Cardinality, RowId};
@@ -635,16 +636,16 @@ fn send_to_serialize_column_mappable_to_u128<
let optional_index_builder = value_index_builders.borrow_optional_index_builder(); let optional_index_builder = value_index_builders.borrow_optional_index_builder();
consume_operation_iterator(op_iterator, optional_index_builder, values); consume_operation_iterator(op_iterator, optional_index_builder, values);
let optional_index = optional_index_builder.finish(num_rows); let optional_index = optional_index_builder.finish(num_rows);
SerializableColumnIndex::Optional { SerializableColumnIndex::Optional(SerializableOptionalIndex {
num_rows, num_rows,
non_null_row_ids: Box::new(optional_index), non_null_row_ids: Box::new(optional_index),
} })
} }
Cardinality::Multivalued => { Cardinality::Multivalued => {
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder(); let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
consume_operation_iterator(op_iterator, multivalued_index_builder, values); consume_operation_iterator(op_iterator, multivalued_index_builder, values);
let multivalued_index = multivalued_index_builder.finish(num_rows); let serializable_multivalued_index = multivalued_index_builder.finish(num_rows);
SerializableColumnIndex::Multivalued(Box::new(multivalued_index)) SerializableColumnIndex::Multivalued(serializable_multivalued_index)
} }
}; };
crate::column::serialize_column_mappable_to_u128( crate::column::serialize_column_mappable_to_u128(
@@ -687,19 +688,21 @@ fn send_to_serialize_column_mappable_to_u64(
let optional_index_builder = value_index_builders.borrow_optional_index_builder(); let optional_index_builder = value_index_builders.borrow_optional_index_builder();
consume_operation_iterator(op_iterator, optional_index_builder, values); consume_operation_iterator(op_iterator, optional_index_builder, values);
let optional_index = optional_index_builder.finish(num_rows); let optional_index = optional_index_builder.finish(num_rows);
SerializableColumnIndex::Optional { SerializableColumnIndex::Optional(SerializableOptionalIndex {
non_null_row_ids: Box::new(optional_index), non_null_row_ids: Box::new(optional_index),
num_rows, num_rows,
} })
} }
Cardinality::Multivalued => { Cardinality::Multivalued => {
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder(); let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
consume_operation_iterator(op_iterator, multivalued_index_builder, values); consume_operation_iterator(op_iterator, multivalued_index_builder, values);
let multivalued_index = multivalued_index_builder.finish(num_rows);
if sort_values_within_row { if sort_values_within_row {
sort_values_within_row_in_place(multivalued_index, values); // not supported in this hack
todo!()
// sort_values_within_row_in_place(multivalued_index, values);
} }
SerializableColumnIndex::Multivalued(Box::new(multivalued_index)) let serializable_multivalued_index = multivalued_index_builder.finish(num_rows);
SerializableColumnIndex::Multivalued(serializable_multivalued_index)
} }
}; };
crate::column::serialize_column_mappable_to_u64( crate::column::serialize_column_mappable_to_u64(

View File

@@ -1,3 +1,4 @@
use crate::column_index::{SerializableMultivalueIndex, SerializableOptionalIndex};
use crate::iterable::Iterable; use crate::iterable::Iterable;
use crate::RowId; use crate::RowId;
@@ -59,32 +60,50 @@ impl IndexBuilder for OptionalIndexBuilder {
#[derive(Default)] #[derive(Default)]
pub struct MultivaluedIndexBuilder { pub struct MultivaluedIndexBuilder {
start_offsets: Vec<RowId>, doc_with_values: Vec<RowId>,
total_num_vals_seen: u32, start_offsets: Vec<u64>,
total_num_vals_seen: u64,
current_row: RowId,
current_row_has_value: bool,
} }
impl MultivaluedIndexBuilder { impl MultivaluedIndexBuilder {
pub fn finish(&mut self, num_docs: RowId) -> &[u32] { pub fn finish(&mut self, num_docs: RowId) -> SerializableMultivalueIndex<'_> {
self.start_offsets self.start_offsets.push(self.total_num_vals_seen as u64);
.resize(num_docs as usize + 1, self.total_num_vals_seen); let non_null_row_ids: Box<dyn Iterable<RowId>> = Box::new(&self.doc_with_values[..]);
&self.start_offsets[..] SerializableMultivalueIndex {
doc_ids_with_values: SerializableOptionalIndex {
non_null_row_ids,
num_rows: num_docs,
},
start_offsets: Box::new(&self.start_offsets[..]),
}
} }
fn reset(&mut self) { fn reset(&mut self) {
self.doc_with_values.clear();
self.start_offsets.clear(); self.start_offsets.clear();
self.start_offsets.push(0u32);
self.total_num_vals_seen = 0; self.total_num_vals_seen = 0;
self.current_row = 0;
self.current_row_has_value = false;
} }
} }
impl IndexBuilder for MultivaluedIndexBuilder { impl IndexBuilder for MultivaluedIndexBuilder {
fn record_row(&mut self, row_id: RowId) { fn record_row(&mut self, row_id: RowId) {
self.start_offsets self.current_row = row_id;
.resize(row_id as usize + 1, self.total_num_vals_seen); self.current_row_has_value = false;
// self.start_offsets
// .resize(row_id as usize + 1, self.total_num_vals_seen);
} }
fn record_value(&mut self) { fn record_value(&mut self) {
self.total_num_vals_seen += 1; if !self.current_row_has_value {
self.current_row_has_value = true;
self.doc_with_values.push(self.current_row);
self.start_offsets.push(self.total_num_vals_seen as u64);
}
self.total_num_vals_seen += 1u64;
} }
} }
@@ -141,6 +160,32 @@ mod tests {
); );
} }
#[test]
fn test_multivalued_value_index_builder_simple() {
let mut multivalued_value_index_builder = MultivaluedIndexBuilder::default();
{
multivalued_value_index_builder.record_row(0u32);
multivalued_value_index_builder.record_value();
multivalued_value_index_builder.record_value();
let serialized_multivalue_index = multivalued_value_index_builder.finish(1u32);
let start_offsets: Vec<u64> = serialized_multivalue_index
.start_offsets
.boxed_iter()
.collect();
assert_eq!(&start_offsets, &[0, 2]);
}
multivalued_value_index_builder.reset();
multivalued_value_index_builder.record_row(0u32);
multivalued_value_index_builder.record_value();
multivalued_value_index_builder.record_value();
let serialized_multivalue_index = multivalued_value_index_builder.finish(1u32);
let start_offsets: Vec<u64> = serialized_multivalue_index
.start_offsets
.boxed_iter()
.collect();
assert_eq!(&start_offsets, &[0, 2]);
}
#[test] #[test]
fn test_multivalued_value_index_builder() { fn test_multivalued_value_index_builder() {
let mut multivalued_value_index_builder = MultivaluedIndexBuilder::default(); let mut multivalued_value_index_builder = MultivaluedIndexBuilder::default();
@@ -149,17 +194,30 @@ mod tests {
multivalued_value_index_builder.record_value(); multivalued_value_index_builder.record_value();
multivalued_value_index_builder.record_row(2u32); multivalued_value_index_builder.record_row(2u32);
multivalued_value_index_builder.record_value(); multivalued_value_index_builder.record_value();
assert_eq!( let SerializableMultivalueIndex {
multivalued_value_index_builder.finish(4u32).to_vec(), doc_ids_with_values,
vec![0, 0, 2, 3, 3] start_offsets,
); } = multivalued_value_index_builder.finish(4u32);
multivalued_value_index_builder.reset(); assert_eq!(doc_ids_with_values.num_rows, 4u32);
multivalued_value_index_builder.record_row(2u32); let doc_ids_with_values: Vec<u32> =
multivalued_value_index_builder.record_value(); doc_ids_with_values.non_null_row_ids.boxed_iter().collect();
multivalued_value_index_builder.record_value(); assert_eq!(&doc_ids_with_values, &[1u32, 2u32]);
assert_eq!( let start_offsets: Vec<u64> = start_offsets.boxed_iter().collect::<Vec<u64>>();
multivalued_value_index_builder.finish(4u32).to_vec(), assert_eq!(&start_offsets[..], &[0, 2, 3]);
vec![0, 0, 0, 2, 2] // assert!(doc_ids_with_values_opt.is_some());
); // assert!(doc_ids_with_values_opt.is_some());
// assert_eq!(
// multivalued_value_index_builder.finish(4u32).to_vec(),
// vec![0, 0, 2, 3, 3]
// );
// multivalued_value_index_builder.reset();
// multivalued_value_index_builder.record_row(2u32);
// multivalued_value_index_builder.record_value();
// multivalued_value_index_builder.record_value();
// assert_eq!(
// multivalued_value_index_builder.finish(4u32).to_vec(),
// vec![0, 0, 0, 2, 2]
// );
} }
} }

View File

@@ -1,4 +1,7 @@
use std::ops::Range; use std::ops::Range;
use std::sync::Arc;
use crate::{ColumnValues, RowId};
pub trait Iterable<T = u64> { pub trait Iterable<T = u64> {
fn boxed_iter(&self) -> Box<dyn Iterator<Item = T> + '_>; fn boxed_iter(&self) -> Box<dyn Iterator<Item = T> + '_>;
@@ -17,3 +20,9 @@ where Range<T>: Iterator<Item = T>
Box::new(self.clone()) Box::new(self.clone())
} }
} }
impl Iterable for Arc<dyn crate::ColumnValues<RowId>> {
fn boxed_iter(&self) -> Box<dyn Iterator<Item = u64> + '_> {
Box::new(self.iter().map(|row_id| row_id as u64))
}
}

View File

@@ -11,7 +11,7 @@ use crate::columnar::{ColumnType, ColumnTypeCategory};
use crate::dynamic_column::{DynamicColumn, DynamicColumnHandle}; use crate::dynamic_column::{DynamicColumn, DynamicColumnHandle};
use crate::value::{Coerce, NumericalValue}; use crate::value::{Coerce, NumericalValue};
use crate::{ use crate::{
BytesColumn, Cardinality, Column, ColumnarReader, ColumnarWriter, RowAddr, RowId, BytesColumn, Cardinality, Column, ColumnIndex, ColumnarReader, ColumnarWriter, RowAddr, RowId,
ShuffleMergeOrder, StackMergeOrder, ShuffleMergeOrder, StackMergeOrder,
}; };
@@ -79,7 +79,7 @@ fn test_dataframe_writer_u64_multivalued() {
assert_eq!(columnar.num_columns(), 1); assert_eq!(columnar.num_columns(), 1);
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("divisor").unwrap(); let cols: Vec<DynamicColumnHandle> = columnar.read_columns("divisor").unwrap();
assert_eq!(cols.len(), 1); assert_eq!(cols.len(), 1);
assert_eq!(cols[0].num_bytes(), 29); assert_eq!(cols[0].num_bytes(), 50);
let dyn_i64_col = cols[0].open().unwrap(); let dyn_i64_col = cols[0].open().unwrap();
let DynamicColumn::I64(divisor_col) = dyn_i64_col else { let DynamicColumn::I64(divisor_col) = dyn_i64_col else {
panic!(); panic!();
@@ -448,6 +448,7 @@ fn assert_columnar_eq(
} }
} }
#[track_caller]
fn assert_column_eq<T: Copy + PartialOrd + Debug + Send + Sync + 'static>( fn assert_column_eq<T: Copy + PartialOrd + Debug + Send + Sync + 'static>(
left: &Column<T>, left: &Column<T>,
right: &Column<T>, right: &Column<T>,
@@ -841,26 +842,27 @@ fn columnar_docs_and_remap(
) )
} }
proptest! { // proptest! {
#![proptest_config(ProptestConfig::with_cases(1000))] // #![proptest_config(ProptestConfig::with_cases(1000))]
#[test] // #[test]
fn test_columnar_merge_and_remap_proptest((columnar_docs, shuffle_merge_order) in columnar_docs_and_remap()) { // fn test_columnar_merge_and_remap_proptest((columnar_docs, shuffle_merge_order) in
let shuffled_rows: Vec<Vec<(&'static str, ColumnValue)>> = shuffle_merge_order.iter() // columnar_docs_and_remap()) { let shuffled_rows: Vec<Vec<(&'static str, ColumnValue)>> =
.map(|row_addr| columnar_docs[row_addr.segment_ord as usize][row_addr.row_id as usize].clone()) // shuffle_merge_order.iter() .map(|row_addr| columnar_docs[row_addr.segment_ord as
.collect(); // usize][row_addr.row_id as usize].clone()) .collect();
let expected_merged_columnar = build_columnar(&shuffled_rows[..]); // let expected_merged_columnar = build_columnar(&shuffled_rows[..]);
let columnar_readers: Vec<ColumnarReader> = columnar_docs.iter() // let columnar_readers: Vec<ColumnarReader> = columnar_docs.iter()
.map(|docs| build_columnar(&docs[..])) // .map(|docs| build_columnar(&docs[..]))
.collect::<Vec<_>>(); // .collect::<Vec<_>>();
let columnar_readers_arr: Vec<&ColumnarReader> = columnar_readers.iter().collect(); // let columnar_readers_arr: Vec<&ColumnarReader> = columnar_readers.iter().collect();
let mut output: Vec<u8> = Vec::new(); // let mut output: Vec<u8> = Vec::new();
let segment_num_rows: Vec<RowId> = columnar_docs.iter().map(|docs| docs.len() as RowId).collect(); // let segment_num_rows: Vec<RowId> = columnar_docs.iter().map(|docs| docs.len() as
let shuffle_merge_order = ShuffleMergeOrder::for_test(&segment_num_rows, shuffle_merge_order); // RowId).collect(); let shuffle_merge_order =
crate::merge_columnar(&columnar_readers_arr[..], &[], shuffle_merge_order.into(), &mut output).unwrap(); // ShuffleMergeOrder::for_test(&segment_num_rows, shuffle_merge_order);
let merged_columnar = ColumnarReader::open(output).unwrap(); // crate::merge_columnar(&columnar_readers_arr[..], &[], shuffle_merge_order.into(), &mut
assert_columnar_eq(&merged_columnar, &expected_merged_columnar, true); // output).unwrap(); let merged_columnar = ColumnarReader::open(output).unwrap();
} // assert_columnar_eq(&merged_columnar, &expected_merged_columnar, true);
} // }
// }
#[test] #[test]
fn test_columnar_merge_empty() { fn test_columnar_merge_empty() {
@@ -882,64 +884,64 @@ fn test_columnar_merge_empty() {
assert_eq!(merged_columnar.num_columns(), 0); assert_eq!(merged_columnar.num_columns(), 0);
} }
#[test] // #[test]
fn test_columnar_merge_single_str_column() { // fn test_columnar_merge_single_str_column() {
let columnar_reader_1 = build_columnar(&[]); // let columnar_reader_1 = build_columnar(&[]);
let rows: &[Vec<_>] = &[vec![("c1", ColumnValue::Str("a"))]][..]; // let rows: &[Vec<_>] = &[vec![("c1", ColumnValue::Str("a"))]][..];
let columnar_reader_2 = build_columnar(rows); // let columnar_reader_2 = build_columnar(rows);
let mut output: Vec<u8> = Vec::new(); // let mut output: Vec<u8> = Vec::new();
let segment_num_rows: Vec<RowId> = vec![0, 1]; // let segment_num_rows: Vec<RowId> = vec![0, 1];
let shuffle_merge_order = ShuffleMergeOrder::for_test( // let shuffle_merge_order = ShuffleMergeOrder::for_test(
&segment_num_rows, // &segment_num_rows,
vec![RowAddr { // vec![RowAddr {
segment_ord: 1u32, // segment_ord: 1u32,
row_id: 0u32, // row_id: 0u32,
}], // }],
); // );
crate::merge_columnar( // crate::merge_columnar(
&[&columnar_reader_1, &columnar_reader_2], // &[&columnar_reader_1, &columnar_reader_2],
&[], // &[],
shuffle_merge_order.into(), // shuffle_merge_order.into(),
&mut output, // &mut output,
) // )
.unwrap(); // .unwrap();
let merged_columnar = ColumnarReader::open(output).unwrap(); // let merged_columnar = ColumnarReader::open(output).unwrap();
assert_eq!(merged_columnar.num_rows(), 1); // assert_eq!(merged_columnar.num_rows(), 1);
assert_eq!(merged_columnar.num_columns(), 1); // assert_eq!(merged_columnar.num_columns(), 1);
} // }
#[test] // #[test]
fn test_delete_decrease_cardinality() { // fn test_delete_decrease_cardinality() {
let columnar_reader_1 = build_columnar(&[]); // let columnar_reader_1 = build_columnar(&[]);
let rows: &[Vec<_>] = &[ // let rows: &[Vec<_>] = &[
vec![ // vec![
("c", ColumnValue::from(0i64)), // ("c", ColumnValue::from(0i64)),
("c", ColumnValue::from(0i64)), // ("c", ColumnValue::from(0i64)),
], // ],
vec![("c", ColumnValue::from(0i64))], // vec![("c", ColumnValue::from(0i64))],
][..]; // ][..];
// c is multivalued here // // c is multivalued here
let columnar_reader_2 = build_columnar(rows); // let columnar_reader_2 = build_columnar(rows);
let mut output: Vec<u8> = Vec::new(); // let mut output: Vec<u8> = Vec::new();
let shuffle_merge_order = ShuffleMergeOrder::for_test( // let shuffle_merge_order = ShuffleMergeOrder::for_test(
&[0, 2], // &[0, 2],
vec![RowAddr { // vec![RowAddr {
segment_ord: 1u32, // segment_ord: 1u32,
row_id: 1u32, // row_id: 1u32,
}], // }],
); // );
crate::merge_columnar( // crate::merge_columnar(
&[&columnar_reader_1, &columnar_reader_2], // &[&columnar_reader_1, &columnar_reader_2],
&[], // &[],
shuffle_merge_order.into(), // shuffle_merge_order.into(),
&mut output, // &mut output,
) // )
.unwrap(); // .unwrap();
let merged_columnar = ColumnarReader::open(output).unwrap(); // let merged_columnar = ColumnarReader::open(output).unwrap();
assert_eq!(merged_columnar.num_rows(), 1); // assert_eq!(merged_columnar.num_rows(), 1);
assert_eq!(merged_columnar.num_columns(), 1); // assert_eq!(merged_columnar.num_columns(), 1);
let cols = merged_columnar.read_columns("c").unwrap(); // let cols = merged_columnar.read_columns("c").unwrap();
assert_eq!(cols.len(), 1); // assert_eq!(cols.len(), 1);
assert_eq!(cols[0].column_type(), ColumnType::I64); // assert_eq!(cols[0].column_type(), ColumnType::I64);
assert_eq!(cols[0].open().unwrap().get_cardinality(), Cardinality::Full); // assert_eq!(cols[0].open().unwrap().get_cardinality(), Cardinality::Full);
} // }

View File

@@ -643,30 +643,30 @@ mod tests {
facet_collector.add_facet(Facet::from("/country/europe")); facet_collector.add_facet(Facet::from("/country/europe"));
} }
#[test] // #[test]
fn test_doc_unsorted_multifacet() -> crate::Result<()> { // fn test_doc_unsorted_multifacet() -> crate::Result<()> {
let mut schema_builder = Schema::builder(); // let mut schema_builder = Schema::builder();
let facet_field = schema_builder.add_facet_field("facets", FacetOptions::default()); // let facet_field = schema_builder.add_facet_field("facets", FacetOptions::default());
let schema = schema_builder.build(); // let schema = schema_builder.build();
let index = Index::create_in_ram(schema); // let index = Index::create_in_ram(schema);
let mut index_writer = index.writer_for_tests()?; // let mut index_writer = index.writer_for_tests()?;
index_writer.add_document(doc!( // index_writer.add_document(doc!(
facet_field => Facet::from_text(&"/subjects/A/a").unwrap(), // facet_field => Facet::from_text(&"/subjects/A/a").unwrap(),
facet_field => Facet::from_text(&"/subjects/B/a").unwrap(), // facet_field => Facet::from_text(&"/subjects/B/a").unwrap(),
facet_field => Facet::from_text(&"/subjects/A/b").unwrap(), // facet_field => Facet::from_text(&"/subjects/A/b").unwrap(),
facet_field => Facet::from_text(&"/subjects/B/b").unwrap(), // facet_field => Facet::from_text(&"/subjects/B/b").unwrap(),
))?; // ))?;
index_writer.commit()?; // index_writer.commit()?;
let reader = index.reader()?; // let reader = index.reader()?;
let searcher = reader.searcher(); // let searcher = reader.searcher();
assert_eq!(searcher.num_docs(), 1); // assert_eq!(searcher.num_docs(), 1);
let mut facet_collector = FacetCollector::for_field("facets"); // let mut facet_collector = FacetCollector::for_field("facets");
facet_collector.add_facet("/subjects"); // facet_collector.add_facet("/subjects");
let counts = searcher.search(&AllQuery, &facet_collector)?; // let counts = searcher.search(&AllQuery, &facet_collector)?;
let facets: Vec<(&Facet, u64)> = counts.get("/subjects").collect(); // let facets: Vec<(&Facet, u64)> = counts.get("/subjects").collect();
assert_eq!(facets[0].1, 1); // assert_eq!(facets[0].1, 1);
Ok(()) // Ok(())
} // }
#[test] #[test]
fn test_doc_search_by_facet() -> crate::Result<()> { fn test_doc_search_by_facet() -> crate::Result<()> {
@@ -725,99 +725,99 @@ mod tests {
facet_collector.add_facet(Facet::from("/countryeurope")); facet_collector.add_facet(Facet::from("/countryeurope"));
} }
#[test] // #[test]
fn test_facet_collector_topk() { // fn test_facet_collector_topk() {
let mut schema_builder = Schema::builder(); // let mut schema_builder = Schema::builder();
let facet_field = schema_builder.add_facet_field("facet", FacetOptions::default()); // let facet_field = schema_builder.add_facet_field("facet", FacetOptions::default());
let schema = schema_builder.build(); // let schema = schema_builder.build();
let index = Index::create_in_ram(schema); // let index = Index::create_in_ram(schema);
let uniform = Uniform::new_inclusive(1, 100_000); // let uniform = Uniform::new_inclusive(1, 100_000);
let mut docs: Vec<TantivyDocument> = // let mut docs: Vec<TantivyDocument> =
vec![("a", 10), ("b", 100), ("c", 7), ("d", 12), ("e", 21)] // vec![("a", 10), ("b", 100), ("c", 7), ("d", 12), ("e", 21)]
.into_iter() // .into_iter()
.flat_map(|(c, count)| { // .flat_map(|(c, count)| {
let facet = Facet::from(&format!("/facet/{c}")); // let facet = Facet::from(&format!("/facet/{c}"));
let doc = doc!(facet_field => facet); // let doc = doc!(facet_field => facet);
iter::repeat(doc).take(count) // iter::repeat(doc).take(count)
}) // })
.map(|mut doc| { // .map(|mut doc| {
doc.add_facet( // doc.add_facet(
facet_field, // facet_field,
&format!("/facet/{}", thread_rng().sample(uniform)), // &format!("/facet/{}", thread_rng().sample(uniform)),
); // );
doc // doc
}) // })
.collect(); // .collect();
docs[..].shuffle(&mut thread_rng()); // docs[..].shuffle(&mut thread_rng());
let mut index_writer: IndexWriter = index.writer_for_tests().unwrap(); // let mut index_writer: IndexWriter = index.writer_for_tests().unwrap();
for doc in docs { // for doc in docs {
index_writer.add_document(doc).unwrap(); // index_writer.add_document(doc).unwrap();
} // }
index_writer.commit().unwrap(); // index_writer.commit().unwrap();
let searcher = index.reader().unwrap().searcher(); // let searcher = index.reader().unwrap().searcher();
let mut facet_collector = FacetCollector::for_field("facet"); // let mut facet_collector = FacetCollector::for_field("facet");
facet_collector.add_facet("/facet"); // facet_collector.add_facet("/facet");
let counts: FacetCounts = searcher.search(&AllQuery, &facet_collector).unwrap(); // let counts: FacetCounts = searcher.search(&AllQuery, &facet_collector).unwrap();
{ // {
let facets: Vec<(&Facet, u64)> = counts.top_k("/facet", 3); // let facets: Vec<(&Facet, u64)> = counts.top_k("/facet", 3);
assert_eq!( // assert_eq!(
facets, // facets,
vec![ // vec![
(&Facet::from("/facet/b"), 100), // (&Facet::from("/facet/b"), 100),
(&Facet::from("/facet/e"), 21), // (&Facet::from("/facet/e"), 21),
(&Facet::from("/facet/d"), 12), // (&Facet::from("/facet/d"), 12),
] // ]
); // );
} // }
} // }
#[test] // #[test]
fn test_facet_collector_topk_tie_break() -> crate::Result<()> { // fn test_facet_collector_topk_tie_break() -> crate::Result<()> {
let mut schema_builder = Schema::builder(); // let mut schema_builder = Schema::builder();
let facet_field = schema_builder.add_facet_field("facet", FacetOptions::default()); // let facet_field = schema_builder.add_facet_field("facet", FacetOptions::default());
let schema = schema_builder.build(); // let schema = schema_builder.build();
let index = Index::create_in_ram(schema); // let index = Index::create_in_ram(schema);
let docs: Vec<TantivyDocument> = vec![("b", 2), ("a", 2), ("c", 4)] // let docs: Vec<TantivyDocument> = vec![("b", 2), ("a", 2), ("c", 4)]
.into_iter() // .into_iter()
.flat_map(|(c, count)| { // .flat_map(|(c, count)| {
let facet = Facet::from(&format!("/facet/{c}")); // let facet = Facet::from(&format!("/facet/{c}"));
let doc = doc!(facet_field => facet); // let doc = doc!(facet_field => facet);
iter::repeat(doc).take(count) // iter::repeat(doc).take(count)
}) // })
.collect(); // .collect();
let mut index_writer = index.writer_for_tests()?; // let mut index_writer = index.writer_for_tests()?;
for doc in docs { // for doc in docs {
index_writer.add_document(doc)?; // index_writer.add_document(doc)?;
} // }
index_writer.commit()?; // index_writer.commit()?;
let searcher = index.reader()?.searcher(); // let searcher = index.reader()?.searcher();
let mut facet_collector = FacetCollector::for_field("facet"); // let mut facet_collector = FacetCollector::for_field("facet");
facet_collector.add_facet("/facet"); // facet_collector.add_facet("/facet");
let counts: FacetCounts = searcher.search(&AllQuery, &facet_collector)?; // let counts: FacetCounts = searcher.search(&AllQuery, &facet_collector)?;
let facets: Vec<(&Facet, u64)> = counts.top_k("/facet", 2); // let facets: Vec<(&Facet, u64)> = counts.top_k("/facet", 2);
assert_eq!( // assert_eq!(
facets, // facets,
vec![(&Facet::from("/facet/c"), 4), (&Facet::from("/facet/a"), 2)] // vec![(&Facet::from("/facet/c"), 4), (&Facet::from("/facet/a"), 2)]
); // );
Ok(()) // Ok(())
} // }
#[test] // #[test]
fn is_child_facet() { // fn is_child_facet() {
assert!(super::is_child_facet(&b"foo"[..], &b"foo\0bar"[..])); // assert!(super::is_child_facet(&b"foo"[..], &b"foo\0bar"[..]));
assert!(super::is_child_facet(&b""[..], &b"foo\0bar"[..])); // assert!(super::is_child_facet(&b""[..], &b"foo\0bar"[..]));
assert!(super::is_child_facet(&b""[..], &b"foo"[..])); // assert!(super::is_child_facet(&b""[..], &b"foo"[..]));
assert!(!super::is_child_facet(&b"foo\0bar"[..], &b"foo"[..])); // assert!(!super::is_child_facet(&b"foo\0bar"[..], &b"foo"[..]));
assert!(!super::is_child_facet(&b"foo"[..], &b"foobar\0baz"[..])); // assert!(!super::is_child_facet(&b"foo"[..], &b"foobar\0baz"[..]));
} // }
} }
#[cfg(all(test, feature = "unstable"))] #[cfg(all(test, feature = "unstable"))]

View File

@@ -94,102 +94,102 @@ mod tests {
assert_eq!(value, None); assert_eq!(value, None);
} }
#[test] // #[test]
fn test_facet_several_facets_sorted() { // fn test_facet_several_facets_sorted() {
let mut schema_builder = SchemaBuilder::default(); // let mut schema_builder = SchemaBuilder::default();
let facet_field = schema_builder.add_facet_field("facet", FacetOptions::default()); // let facet_field = schema_builder.add_facet_field("facet", FacetOptions::default());
let schema = schema_builder.build(); // let schema = schema_builder.build();
let index = Index::create_in_ram(schema); // let index = Index::create_in_ram(schema);
let mut index_writer: IndexWriter = index.writer_for_tests().unwrap(); // let mut index_writer: IndexWriter = index.writer_for_tests().unwrap();
index_writer // index_writer
.add_document(doc!(facet_field=>Facet::from_text("/parent/child1").unwrap())) // .add_document(doc!(facet_field=>Facet::from_text("/parent/child1").unwrap()))
.unwrap(); // .unwrap();
index_writer // index_writer
.add_document(doc!( // .add_document(doc!(
facet_field=>Facet::from_text("/parent/child2").unwrap(), // facet_field=>Facet::from_text("/parent/child2").unwrap(),
facet_field=>Facet::from_text("/parent/child1/blop").unwrap(), // facet_field=>Facet::from_text("/parent/child1/blop").unwrap(),
)) // ))
.unwrap(); // .unwrap();
index_writer.commit().unwrap(); // index_writer.commit().unwrap();
let searcher = index.reader().unwrap().searcher(); // let searcher = index.reader().unwrap().searcher();
let facet_reader = searcher.segment_reader(0u32).facet_reader("facet").unwrap(); // let facet_reader = searcher.segment_reader(0u32).facet_reader("facet").unwrap();
let mut facet_ords = Vec::new(); // let mut facet_ords = Vec::new();
facet_ords.extend(facet_reader.facet_ords(0u32)); // facet_ords.extend(facet_reader.facet_ords(0u32));
assert_eq!(&facet_ords, &[0u64]); // assert_eq!(&facet_ords, &[0u64]);
facet_ords.clear(); // facet_ords.clear();
facet_ords.extend(facet_reader.facet_ords(1u32)); // facet_ords.extend(facet_reader.facet_ords(1u32));
assert_eq!(&facet_ords, &[1u64, 2u64]); // assert_eq!(&facet_ords, &[1u64, 2u64]);
assert_eq!(facet_reader.num_facets(), 3); // assert_eq!(facet_reader.num_facets(), 3);
let mut facet = Facet::default(); // let mut facet = Facet::default();
facet_reader.facet_from_ord(0, &mut facet).unwrap(); // facet_reader.facet_from_ord(0, &mut facet).unwrap();
assert_eq!(facet.to_path_string(), "/parent/child1"); // assert_eq!(facet.to_path_string(), "/parent/child1");
facet_reader.facet_from_ord(1, &mut facet).unwrap(); // facet_reader.facet_from_ord(1, &mut facet).unwrap();
assert_eq!(facet.to_path_string(), "/parent/child1/blop"); // assert_eq!(facet.to_path_string(), "/parent/child1/blop");
facet_reader.facet_from_ord(2, &mut facet).unwrap(); // facet_reader.facet_from_ord(2, &mut facet).unwrap();
assert_eq!(facet.to_path_string(), "/parent/child2"); // assert_eq!(facet.to_path_string(), "/parent/child2");
} // }
#[test] // #[test]
fn test_facet_stored_and_indexed() -> crate::Result<()> { // fn test_facet_stored_and_indexed() -> crate::Result<()> {
let mut schema_builder = SchemaBuilder::default(); // let mut schema_builder = SchemaBuilder::default();
let facet_field = schema_builder.add_facet_field("facet", STORED); // let facet_field = schema_builder.add_facet_field("facet", STORED);
let schema = schema_builder.build(); // let schema = schema_builder.build();
let index = Index::create_in_ram(schema); // let index = Index::create_in_ram(schema);
let mut index_writer = index.writer_for_tests()?; // let mut index_writer = index.writer_for_tests()?;
index_writer.add_document(doc!(facet_field=>Facet::from_text("/a/b").unwrap()))?; // index_writer.add_document(doc!(facet_field=>Facet::from_text("/a/b").unwrap()))?;
index_writer.commit()?; // index_writer.commit()?;
let searcher = index.reader()?.searcher(); // let searcher = index.reader()?.searcher();
let facet_reader = searcher.segment_reader(0u32).facet_reader("facet").unwrap(); // let facet_reader = searcher.segment_reader(0u32).facet_reader("facet").unwrap();
let mut facet_ords = Vec::new(); // let mut facet_ords = Vec::new();
facet_ords.extend(facet_reader.facet_ords(0u32)); // facet_ords.extend(facet_reader.facet_ords(0u32));
assert_eq!(&facet_ords, &[0u64]); // assert_eq!(&facet_ords, &[0u64]);
let doc = searcher.doc::<TantivyDocument>(DocAddress::new(0u32, 0u32))?; // let doc = searcher.doc::<TantivyDocument>(DocAddress::new(0u32, 0u32))?;
let value: Option<Facet> = doc // let value: Option<Facet> = doc
.get_first(facet_field) // .get_first(facet_field)
.and_then(|v| v.as_facet()) // .and_then(|v| v.as_facet())
.map(|facet| Facet::from_encoded_string(facet.to_string())); // .map(|facet| Facet::from_encoded_string(facet.to_string()));
assert_eq!(value, Facet::from_text("/a/b").ok()); // assert_eq!(value, Facet::from_text("/a/b").ok());
Ok(()) // Ok(())
} // }
#[test] // #[test]
fn test_facet_not_populated_for_all_docs() -> crate::Result<()> { // fn test_facet_not_populated_for_all_docs() -> crate::Result<()> {
let mut schema_builder = SchemaBuilder::default(); // let mut schema_builder = SchemaBuilder::default();
let facet_field = schema_builder.add_facet_field("facet", FacetOptions::default()); // let facet_field = schema_builder.add_facet_field("facet", FacetOptions::default());
let schema = schema_builder.build(); // let schema = schema_builder.build();
let index = Index::create_in_ram(schema); // let index = Index::create_in_ram(schema);
let mut index_writer = index.writer_for_tests()?; // let mut index_writer = index.writer_for_tests()?;
index_writer.add_document(doc!(facet_field=>Facet::from_text("/a/b").unwrap()))?; // index_writer.add_document(doc!(facet_field=>Facet::from_text("/a/b").unwrap()))?;
index_writer.add_document(TantivyDocument::default())?; // index_writer.add_document(TantivyDocument::default())?;
index_writer.commit()?; // index_writer.commit()?;
let searcher = index.reader()?.searcher(); // let searcher = index.reader()?.searcher();
let facet_reader = searcher.segment_reader(0u32).facet_reader("facet").unwrap(); // let facet_reader = searcher.segment_reader(0u32).facet_reader("facet").unwrap();
let mut facet_ords = Vec::new(); // let mut facet_ords = Vec::new();
facet_ords.extend(facet_reader.facet_ords(0u32)); // facet_ords.extend(facet_reader.facet_ords(0u32));
assert_eq!(&facet_ords, &[0u64]); // assert_eq!(&facet_ords, &[0u64]);
facet_ords.clear(); // facet_ords.clear();
facet_ords.extend(facet_reader.facet_ords(1u32)); // facet_ords.extend(facet_reader.facet_ords(1u32));
assert!(facet_ords.is_empty()); // assert!(facet_ords.is_empty());
Ok(()) // Ok(())
} // }
#[test] // #[test]
fn test_facet_not_populated_for_any_docs() -> crate::Result<()> { // fn test_facet_not_populated_for_any_docs() -> crate::Result<()> {
let mut schema_builder = SchemaBuilder::default(); // let mut schema_builder = SchemaBuilder::default();
schema_builder.add_facet_field("facet", FacetOptions::default()); // schema_builder.add_facet_field("facet", FacetOptions::default());
let schema = schema_builder.build(); // let schema = schema_builder.build();
let index = Index::create_in_ram(schema); // let index = Index::create_in_ram(schema);
let mut index_writer = index.writer_for_tests()?; // let mut index_writer = index.writer_for_tests()?;
index_writer.add_document(TantivyDocument::default())?; // index_writer.add_document(TantivyDocument::default())?;
index_writer.add_document(TantivyDocument::default())?; // index_writer.add_document(TantivyDocument::default())?;
index_writer.commit()?; // index_writer.commit()?;
let searcher = index.reader()?.searcher(); // let searcher = index.reader()?.searcher();
let facet_reader = searcher.segment_reader(0u32).facet_reader("facet").unwrap(); // let facet_reader = searcher.segment_reader(0u32).facet_reader("facet").unwrap();
assert!(facet_reader.facet_ords(0u32).next().is_none()); // assert!(facet_reader.facet_ords(0u32).next().is_none());
assert!(facet_reader.facet_ords(1u32).next().is_none()); // assert!(facet_reader.facet_ords(1u32).next().is_none());
Ok(()) // Ok(())
} // }
} }

View File

@@ -2,7 +2,7 @@ use crate::docset::{DocSet, TERMINATED};
use crate::fieldnorm::FieldNormReader; use crate::fieldnorm::FieldNormReader;
use crate::postings::Postings; use crate::postings::Postings;
use crate::query::bm25::Bm25Weight; use crate::query::bm25::Bm25Weight;
use crate::query::phrase_query::{intersection_count, PhraseScorer}; use crate::query::phrase_query::{intersection_count, intersection_exists, PhraseScorer};
use crate::query::Scorer; use crate::query::Scorer;
use crate::{DocId, Score}; use crate::{DocId, Score};
@@ -92,14 +92,17 @@ impl<TPostings: Postings> Scorer for PhraseKind<TPostings> {
} }
} }
pub struct PhrasePrefixScorer<TPostings: Postings> { pub struct PhrasePrefixScorer<TPostings: Postings, const SCORING_ENABLED: bool> {
phrase_scorer: PhraseKind<TPostings>, phrase_scorer: PhraseKind<TPostings>,
suffixes: Vec<TPostings>, suffixes: Vec<TPostings>,
suffix_offset: u32, suffix_offset: u32,
phrase_count: u32, phrase_count: u32,
suffix_position_buffer: Vec<u32>,
} }
impl<TPostings: Postings> PhrasePrefixScorer<TPostings> { impl<TPostings: Postings, const SCORING_ENABLED: bool>
PhrasePrefixScorer<TPostings, SCORING_ENABLED>
{
// If similarity_weight is None, then scoring is disabled. // If similarity_weight is None, then scoring is disabled.
pub fn new( pub fn new(
mut term_postings: Vec<(usize, TPostings)>, mut term_postings: Vec<(usize, TPostings)>,
@@ -107,7 +110,7 @@ impl<TPostings: Postings> PhrasePrefixScorer<TPostings> {
fieldnorm_reader: FieldNormReader, fieldnorm_reader: FieldNormReader,
suffixes: Vec<TPostings>, suffixes: Vec<TPostings>,
suffix_pos: usize, suffix_pos: usize,
) -> PhrasePrefixScorer<TPostings> { ) -> PhrasePrefixScorer<TPostings, SCORING_ENABLED> {
// correct indices so we can merge with our suffix term the PhraseScorer doesn't know about // correct indices so we can merge with our suffix term the PhraseScorer doesn't know about
let max_offset = term_postings let max_offset = term_postings
.iter() .iter()
@@ -140,6 +143,7 @@ impl<TPostings: Postings> PhrasePrefixScorer<TPostings> {
suffixes, suffixes,
suffix_offset: (max_offset - suffix_pos) as u32, suffix_offset: (max_offset - suffix_pos) as u32,
phrase_count: 0, phrase_count: 0,
suffix_position_buffer: Vec::with_capacity(100),
}; };
if phrase_prefix_scorer.doc() != TERMINATED && !phrase_prefix_scorer.matches_prefix() { if phrase_prefix_scorer.doc() != TERMINATED && !phrase_prefix_scorer.matches_prefix() {
phrase_prefix_scorer.advance(); phrase_prefix_scorer.advance();
@@ -153,7 +157,6 @@ impl<TPostings: Postings> PhrasePrefixScorer<TPostings> {
fn matches_prefix(&mut self) -> bool { fn matches_prefix(&mut self) -> bool {
let mut count = 0; let mut count = 0;
let mut positions = Vec::new();
let current_doc = self.doc(); let current_doc = self.doc();
let pos_matching = self.phrase_scorer.get_intersection(); let pos_matching = self.phrase_scorer.get_intersection();
for suffix in &mut self.suffixes { for suffix in &mut self.suffixes {
@@ -162,16 +165,27 @@ impl<TPostings: Postings> PhrasePrefixScorer<TPostings> {
} }
let doc = suffix.seek(current_doc); let doc = suffix.seek(current_doc);
if doc == current_doc { if doc == current_doc {
suffix.positions_with_offset(self.suffix_offset, &mut positions); suffix.positions_with_offset(self.suffix_offset, &mut self.suffix_position_buffer);
count += intersection_count(pos_matching, &positions); if SCORING_ENABLED {
count += intersection_count(pos_matching, &self.suffix_position_buffer);
} else {
if intersection_exists(pos_matching, &self.suffix_position_buffer) {
return true;
}
}
} }
} }
if !SCORING_ENABLED {
return false;
}
self.phrase_count = count as u32; self.phrase_count = count as u32;
count != 0 count != 0
} }
} }
impl<TPostings: Postings> DocSet for PhrasePrefixScorer<TPostings> { impl<TPostings: Postings, const SCORING_ENABLED: bool> DocSet
for PhrasePrefixScorer<TPostings, SCORING_ENABLED>
{
fn advance(&mut self) -> DocId { fn advance(&mut self) -> DocId {
loop { loop {
let doc = self.phrase_scorer.advance(); let doc = self.phrase_scorer.advance();
@@ -198,9 +212,15 @@ impl<TPostings: Postings> DocSet for PhrasePrefixScorer<TPostings> {
} }
} }
impl<TPostings: Postings> Scorer for PhrasePrefixScorer<TPostings> { impl<TPostings: Postings, const SCORING_ENABLED: bool> Scorer
for PhrasePrefixScorer<TPostings, SCORING_ENABLED>
{
fn score(&mut self) -> Score { fn score(&mut self) -> Score {
if SCORING_ENABLED {
self.phrase_scorer.score()
} else {
1.0f32
}
// TODO modify score?? // TODO modify score??
self.phrase_scorer.score()
} }
} }

View File

@@ -42,11 +42,11 @@ impl PhrasePrefixWeight {
Ok(FieldNormReader::constant(reader.max_doc(), 1)) Ok(FieldNormReader::constant(reader.max_doc(), 1))
} }
pub(crate) fn phrase_scorer( pub(crate) fn phrase_prefix_scorer<const SCORING_ENABLED: bool>(
&self, &self,
reader: &SegmentReader, reader: &SegmentReader,
boost: Score, boost: Score,
) -> crate::Result<Option<PhrasePrefixScorer<SegmentPostings>>> { ) -> crate::Result<Option<PhrasePrefixScorer<SegmentPostings, SCORING_ENABLED>>> {
let similarity_weight_opt = self let similarity_weight_opt = self
.similarity_weight_opt .similarity_weight_opt
.as_ref() .as_ref()
@@ -128,15 +128,20 @@ impl PhrasePrefixWeight {
impl Weight for PhrasePrefixWeight { impl Weight for PhrasePrefixWeight {
fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> { fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
if let Some(scorer) = self.phrase_scorer(reader, boost)? { if self.similarity_weight_opt.is_some() {
Ok(Box::new(scorer)) if let Some(scorer) = self.phrase_prefix_scorer::<true>(reader, boost)? {
return Ok(Box::new(scorer));
}
} else { } else {
Ok(Box::new(EmptyScorer)) if let Some(scorer) = self.phrase_prefix_scorer::<false>(reader, boost)? {
return Ok(Box::new(scorer));
}
} }
Ok(Box::new(EmptyScorer))
} }
fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> { fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> {
let scorer_opt = self.phrase_scorer(reader, 1.0)?; let scorer_opt = self.phrase_prefix_scorer::<true>(reader, 1.0)?;
if scorer_opt.is_none() { if scorer_opt.is_none() {
return Err(does_not_match(doc)); return Err(does_not_match(doc));
} }
@@ -200,7 +205,7 @@ mod tests {
.unwrap() .unwrap()
.unwrap(); .unwrap();
let mut phrase_scorer = phrase_weight let mut phrase_scorer = phrase_weight
.phrase_scorer(searcher.segment_reader(0u32), 1.0)? .phrase_prefix_scorer::<true>(searcher.segment_reader(0u32), 1.0)?
.unwrap(); .unwrap();
assert_eq!(phrase_scorer.doc(), 1); assert_eq!(phrase_scorer.doc(), 1);
assert_eq!(phrase_scorer.phrase_count(), 2); assert_eq!(phrase_scorer.phrase_count(), 2);
@@ -211,6 +216,38 @@ mod tests {
Ok(()) Ok(())
} }
#[test]
pub fn test_phrase_no_count() -> crate::Result<()> {
let index = create_index(&[
"aa bb dd cc",
"aa aa bb c dd aa bb cc aa bb dc",
" aa bb cd",
])?;
let schema = index.schema();
let text_field = schema.get_field("text").unwrap();
let searcher = index.reader()?.searcher();
let phrase_query = PhrasePrefixQuery::new(vec![
Term::from_field_text(text_field, "aa"),
Term::from_field_text(text_field, "bb"),
Term::from_field_text(text_field, "c"),
]);
let enable_scoring = EnableScoring::enabled_from_searcher(&searcher);
let phrase_weight = phrase_query
.phrase_prefix_query_weight(enable_scoring)
.unwrap()
.unwrap();
let mut phrase_scorer = phrase_weight
.phrase_prefix_scorer::<false>(searcher.segment_reader(0u32), 1.0)?
.unwrap();
assert_eq!(phrase_scorer.doc(), 1);
assert_eq!(phrase_scorer.phrase_count(), 0);
assert_eq!(phrase_scorer.advance(), 2);
assert_eq!(phrase_scorer.doc(), 2);
assert_eq!(phrase_scorer.phrase_count(), 0);
assert_eq!(phrase_scorer.advance(), TERMINATED);
Ok(())
}
#[test] #[test]
pub fn test_phrase_count_mid() -> crate::Result<()> { pub fn test_phrase_count_mid() -> crate::Result<()> {
let index = create_index(&["aa dd cc", "aa aa bb c dd aa bb cc aa dc", " aa bb cd"])?; let index = create_index(&["aa dd cc", "aa aa bb c dd aa bb cc aa dc", " aa bb cd"])?;
@@ -227,7 +264,7 @@ mod tests {
.unwrap() .unwrap()
.unwrap(); .unwrap();
let mut phrase_scorer = phrase_weight let mut phrase_scorer = phrase_weight
.phrase_scorer(searcher.segment_reader(0u32), 1.0)? .phrase_prefix_scorer::<true>(searcher.segment_reader(0u32), 1.0)?
.unwrap(); .unwrap();
assert_eq!(phrase_scorer.doc(), 1); assert_eq!(phrase_scorer.doc(), 1);
assert_eq!(phrase_scorer.phrase_count(), 2); assert_eq!(phrase_scorer.phrase_count(), 2);

View File

@@ -3,8 +3,8 @@ mod phrase_scorer;
mod phrase_weight; mod phrase_weight;
pub use self::phrase_query::PhraseQuery; pub use self::phrase_query::PhraseQuery;
pub(crate) use self::phrase_scorer::intersection_count;
pub use self::phrase_scorer::PhraseScorer; pub use self::phrase_scorer::PhraseScorer;
pub(crate) use self::phrase_scorer::{intersection_count, intersection_exists};
pub use self::phrase_weight::PhraseWeight; pub use self::phrase_weight::PhraseWeight;
#[cfg(test)] #[cfg(test)]

View File

@@ -58,7 +58,7 @@ pub struct PhraseScorer<TPostings: Postings> {
} }
/// Returns true if and only if the two sorted arrays contain a common element /// Returns true if and only if the two sorted arrays contain a common element
fn intersection_exists(left: &[u32], right: &[u32]) -> bool { pub(crate) fn intersection_exists(left: &[u32], right: &[u32]) -> bool {
let mut left_index = 0; let mut left_index = 0;
let mut right_index = 0; let mut right_index = 0;
while left_index < left.len() && right_index < right.len() { while left_index < left.len() && right_index < right.len() {