Skip to content

Commit

Permalink
Fix bug that can cause get_docids_for_value_range to panic. (#2295)
Browse files Browse the repository at this point in the history
* Fix bug that can cause `get_docids_for_value_range` to panic.

When `selected_docid_range.end == num_rows`, we would get a panic
as we try to access a non-existing blockmeta.

This PR accepts calls to rank with any value.
For any value above num_rows we simply return non_null_rows.

Fixes #2293

* add tests, merge variables

---------

Co-authored-by: Pascal Seitz <[email protected]>
  • Loading branch information
fulmicoton and PSeitz authored Jan 9, 2024
1 parent 53f2fe1 commit 014328e
Show file tree
Hide file tree
Showing 5 changed files with 48 additions and 18 deletions.
12 changes: 6 additions & 6 deletions columnar/src/column_index/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,18 +126,18 @@ impl ColumnIndex {
}
}

pub fn docid_range_to_rowids(&self, doc_id: Range<DocId>) -> Range<RowId> {
pub fn docid_range_to_rowids(&self, doc_id_range: Range<DocId>) -> Range<RowId> {
match self {
ColumnIndex::Empty { .. } => 0..0,
ColumnIndex::Full => doc_id,
ColumnIndex::Full => doc_id_range,
ColumnIndex::Optional(optional_index) => {
let row_start = optional_index.rank(doc_id.start);
let row_end = optional_index.rank(doc_id.end);
let row_start = optional_index.rank(doc_id_range.start);
let row_end = optional_index.rank(doc_id_range.end);
row_start..row_end
}
ColumnIndex::Multivalued(multivalued_index) => {
let end_docid = doc_id.end.min(multivalued_index.num_docs() - 1) + 1;
let start_docid = doc_id.start.min(end_docid);
let end_docid = doc_id_range.end.min(multivalued_index.num_docs() - 1) + 1;
let start_docid = doc_id_range.start.min(end_docid);

let row_start = multivalued_index.start_index_column.get_val(start_docid);
let row_end = multivalued_index.start_index_column.get_val(end_docid);
Expand Down
17 changes: 11 additions & 6 deletions columnar/src/column_index/optional_index/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,6 @@ const DENSE_BLOCK_THRESHOLD: u32 =

const ELEMENTS_PER_BLOCK: u32 = u16::MAX as u32 + 1;

const BLOCK_SIZE: RowId = 1 << 16;

#[derive(Copy, Clone, Debug)]
struct BlockMeta {
non_null_rows_before_block: u32,
Expand Down Expand Up @@ -109,8 +107,8 @@ struct RowAddr {
#[inline(always)]
fn row_addr_from_row_id(row_id: RowId) -> RowAddr {
RowAddr {
block_id: (row_id / BLOCK_SIZE) as u16,
in_block_row_id: (row_id % BLOCK_SIZE) as u16,
block_id: (row_id / ELEMENTS_PER_BLOCK) as u16,
in_block_row_id: (row_id % ELEMENTS_PER_BLOCK) as u16,
}
}

Expand Down Expand Up @@ -185,8 +183,13 @@ impl Set<RowId> for OptionalIndex {
}
}

/// Any value doc_id is allowed.
/// In particular, doc_id = num_rows.
#[inline]
fn rank(&self, doc_id: DocId) -> RowId {
if doc_id >= self.num_docs() {
return self.num_non_nulls();
}
let RowAddr {
block_id,
in_block_row_id,
Expand All @@ -200,13 +203,15 @@ impl Set<RowId> for OptionalIndex {
block_meta.non_null_rows_before_block + block_offset_row_id
}

/// Any value doc_id is allowed.
/// In particular, doc_id = num_rows.
#[inline]
fn rank_if_exists(&self, doc_id: DocId) -> Option<RowId> {
let RowAddr {
block_id,
in_block_row_id,
} = row_addr_from_row_id(doc_id);
let block_meta = self.block_metas[block_id as usize];
let block_meta = *self.block_metas.get(block_id as usize)?;
let block = self.block(block_meta);
let block_offset_row_id = match block {
Block::Dense(dense_block) => dense_block.rank_if_exists(in_block_row_id),
Expand Down Expand Up @@ -491,7 +496,7 @@ fn deserialize_optional_index_block_metadatas(
non_null_rows_before_block += num_non_null_rows;
}
block_metas.resize(
((num_rows + BLOCK_SIZE - 1) / BLOCK_SIZE) as usize,
((num_rows + ELEMENTS_PER_BLOCK - 1) / ELEMENTS_PER_BLOCK) as usize,
BlockMeta {
non_null_rows_before_block,
start_byte_offset,
Expand Down
3 changes: 2 additions & 1 deletion columnar/src/column_index/optional_index/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ pub trait Set<T> {
///
/// # Panics
///
/// May panic if rank is greater than the number of elements in the Set.
/// May panic if rank is greater or equal to the number of
/// elements in the Set.
fn select(&self, rank: T) -> T;

/// Creates a brand new select cursor.
Expand Down
32 changes: 28 additions & 4 deletions columnar/src/column_index/optional_index/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,30 @@ use proptest::strategy::Strategy;
use proptest::{prop_oneof, proptest};

use super::*;
use crate::{ColumnarReader, ColumnarWriter, DynamicColumnHandle};

#[test]
fn test_optional_index_bug_2293() {
// tests for panic in docid_range_to_rowids for docid == num_docs
test_optional_index_with_num_docs(ELEMENTS_PER_BLOCK - 1);
test_optional_index_with_num_docs(ELEMENTS_PER_BLOCK);
test_optional_index_with_num_docs(ELEMENTS_PER_BLOCK + 1);
}
fn test_optional_index_with_num_docs(num_docs: u32) {
let mut dataframe_writer = ColumnarWriter::default();
dataframe_writer.record_numerical(100, "score", 80i64);
let mut buffer: Vec<u8> = Vec::new();
dataframe_writer
.serialize(num_docs, None, &mut buffer)
.unwrap();
let columnar = ColumnarReader::open(buffer).unwrap();
assert_eq!(columnar.num_columns(), 1);
let cols: Vec<DynamicColumnHandle> = columnar.read_columns("score").unwrap();
assert_eq!(cols.len(), 1);

let col = cols[0].open().unwrap();
col.column_index().docid_range_to_rowids(0..num_docs);
}

#[test]
fn test_dense_block_threshold() {
Expand Down Expand Up @@ -35,7 +59,7 @@ proptest! {

#[test]
fn test_with_random_sets_simple() {
let vals = 10..BLOCK_SIZE * 2;
let vals = 10..ELEMENTS_PER_BLOCK * 2;
let mut out: Vec<u8> = Vec::new();
serialize_optional_index(&vals, 100, &mut out).unwrap();
let null_index = open_optional_index(OwnedBytes::new(out)).unwrap();
Expand Down Expand Up @@ -171,7 +195,7 @@ fn test_optional_index_rank() {
test_optional_index_rank_aux(&[0u32, 1u32]);
let mut block = Vec::new();
block.push(3u32);
block.extend((0..BLOCK_SIZE).map(|i| i + BLOCK_SIZE + 1));
block.extend((0..ELEMENTS_PER_BLOCK).map(|i| i + ELEMENTS_PER_BLOCK + 1));
test_optional_index_rank_aux(&block);
}

Expand All @@ -185,8 +209,8 @@ fn test_optional_index_iter_empty_one() {
fn test_optional_index_iter_dense_block() {
let mut block = Vec::new();
block.push(3u32);
block.extend((0..BLOCK_SIZE).map(|i| i + BLOCK_SIZE + 1));
test_optional_index_iter_aux(&block, 3 * BLOCK_SIZE);
block.extend((0..ELEMENTS_PER_BLOCK).map(|i| i + ELEMENTS_PER_BLOCK + 1));
test_optional_index_iter_aux(&block, 3 * ELEMENTS_PER_BLOCK);
}

#[test]
Expand Down
2 changes: 1 addition & 1 deletion columnar/src/column_values/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ pub trait ColumnValues<T: PartialOrd = u64>: Send + Sync {
row_id_hits: &mut Vec<RowId>,
) {
let row_id_range = row_id_range.start..row_id_range.end.min(self.num_vals());
for idx in row_id_range.start..row_id_range.end {
for idx in row_id_range {
let val = self.get_val(idx);
if value_range.contains(&val) {
row_id_hits.push(idx);
Expand Down

0 comments on commit 014328e

Please sign in to comment.