rust/tantivy
1
0
Fork 0
mirror of https://github.com/quickwit-oss/tantivy.git synced 2026-10-07 12:22:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added support for writing plain strings

Browse Source
This commit is contained in:
Paul Masurel
2026-09-15 12:40:36 +02:00
parent c1d0da766e
commit 082074040a
16 changed files with 1553 additions and 826 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
+1 -1
View File
@@ -217,7 +217,7 @@ name = "bool_queries_with_range"
harness = false
[[bench]]
name = "str_search_and_get"
name = "dictionary_encoded_str_search_and_get"
harness = false
[[bench]]
benches/str_search_and_get.rs → benches/dictionary_encoded_str_search_and_get.rs
View File
columnar/src/column/bytes_or_str.rs
+82 -2
View File
@@ -1,9 +1,66 @@
use std::fmt;
use std::sync::Arc;
use std::{fmt, io};
use common::file_slice::FileSlice;
use common::{HasLen, OwnedBytes};
use super::dictionary_encoding::open_dictionary_bytes_column;
use super::plain::open_plain_bytes_column;
use super::{
DictionaryEncodedBytesColumn, DictionaryEncodedStrColumn, PlainBytesColumn, PlainStrColumn,
};
use crate::{Cardinality, ColumnIndex, PayloadEncoding, RowId};
use crate::{Cardinality, ColumnIndex, PayloadEncoding, RowId, Version};
pub fn open_column_bytes(data: OwnedBytes, format_version: Version) -> io::Result<BytesColumn> {
open_column_bytes_from_file_slice(FileSlice::new(Arc::new(data)), format_version)
}
pub(crate) fn open_column_bytes_from_file_slice(
data: FileSlice,
format_version: Version,
) -> io::Result<BytesColumn> {
match format_version {
Version::V1 | Version::V2 => {
open_dictionary_bytes_column(data.read_bytes()?, format_version).map(Into::into)
}
Version::V3 => {
if data.len() < 1 {
return Err(invalid_data("missing string/byte payload encoding tag"));
}
let (encoding_slice, payload) = data.split(1);
let encoding_bytes = encoding_slice.read_bytes()?;
let encoding =
PayloadEncoding::try_from_code(encoding_bytes[0]).map_err(io::Error::from)?;
match encoding {
PayloadEncoding::Dictionary => {
open_dictionary_bytes_column(payload.read_bytes()?, format_version)
.map(Into::into)
}
PayloadEncoding::Plain => open_plain_bytes_column(payload).map(Into::into),
}
}
}
}
pub fn open_column_str(data: OwnedBytes, format_version: Version) -> io::Result<StrColumn> {
open_column_str_from_file_slice(FileSlice::new(Arc::new(data)), format_version)
}
pub(crate) fn open_column_str_from_file_slice(
data: FileSlice,
format_version: Version,
) -> io::Result<StrColumn> {
match open_column_bytes_from_file_slice(data, format_version)? {
BytesColumn::DictionaryEncoded(bytes_column) => {
Ok(DictionaryEncodedStrColumn::wrap(bytes_column).into())
}
BytesColumn::Plain(bytes_column) => Ok(PlainStrColumn::wrap(bytes_column).into()),
}
}
fn invalid_data(message: &'static str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message)
}
/// A byte column, independently of its payload encoding.
#[derive(Clone)]
@@ -207,6 +264,29 @@ impl From<StrColumn> for BytesColumn {
mod tests {
use super::*;
#[test]
fn test_v3_payload_encoding_tag_errors() {
let error = open_column_bytes(OwnedBytes::new(Vec::new()), Version::V3).unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
let error = open_column_bytes(OwnedBytes::new(vec![u8::MAX]), Version::V3).unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
let error = open_column_bytes(
OwnedBytes::new(vec![PayloadEncoding::Plain.to_code()]),
Version::V3,
)
.unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
let error = open_column_bytes(
OwnedBytes::new(vec![PayloadEncoding::Dictionary.to_code()]),
Version::V3,
)
.unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn test_dictionary_encoded_column_metadata_and_downcasts() {
let bytes_column: BytesColumn = DictionaryEncodedBytesColumn::empty(3).into();
columnar/src/column/dictionary_encoded.rs → columnar/src/column/dictionary_encoding/mod.rs
+5 -1
View File
@@ -7,6 +7,10 @@ use sstable::{Dictionary, VoidSSTable};
use crate::column::Column;
use crate::{ColumnIndex, RowId};
mod open;
pub(crate) use open::open_dictionary_bytes_column;
/// Dictionary encoded column.
///
/// The column simply gives access to a regular u64-column that, in
@@ -22,7 +26,7 @@ pub struct DictionaryEncodedBytesColumn {
}
impl fmt::Debug for DictionaryEncodedBytesColumn {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("DictionaryEncodedBytesColumn")
.field("term_ord_column", &self.term_ord_column)
.finish()
columnar/src/column/dictionary_encoding/open.rs
+35
View File
@@ -0,0 +1,35 @@
use std::io;
use std::sync::Arc;
use common::OwnedBytes;
use sstable::Dictionary;
use super::DictionaryEncodedBytesColumn;
use crate::Version;
pub(crate) fn open_dictionary_bytes_column(
data: OwnedBytes,
format_version: Version,
) -> io::Result<DictionaryEncodedBytesColumn> {
if data.len() < 4 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"truncated dictionary string/byte column payload",
));
}
let (body, dictionary_len_bytes) = data.rsplit(4);
let dictionary_len = u32::from_le_bytes(dictionary_len_bytes.as_slice().try_into().unwrap());
if dictionary_len as usize > body.len() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"dictionary length exceeds string/byte column payload",
));
}
let (dictionary_bytes, column_bytes) = body.split(dictionary_len as usize);
let dictionary = Arc::new(Dictionary::from_bytes(dictionary_bytes)?);
let term_ord_column = crate::column::open_column_u64::<u64>(column_bytes, format_version)?;
Ok(DictionaryEncodedBytesColumn {
dictionary,
term_ord_column,
})
}
columnar/src/column/mod.rs
+10 -13
View File
@@ -1,30 +1,27 @@
mod bytes_or_str;
mod dictionary_encoded;
mod dictionary_encoding;
mod numerical;
mod plain;
mod serialize;
use std::fmt::{self, Debug};
use std::io::Write;
use std::ops::{Range, RangeInclusive};
use std::sync::Arc;
pub use bytes_or_str::{BytesColumn, StrColumn};
pub use bytes_or_str::{BytesColumn, StrColumn, open_column_bytes, open_column_str};
pub(crate) use bytes_or_str::{open_column_bytes_from_file_slice, open_column_str_from_file_slice};
use common::BinarySerializable;
pub use dictionary_encoded::{DictionaryEncodedBytesColumn, DictionaryEncodedStrColumn};
pub(crate) use plain::PlainBlockMeta;
#[cfg(test)]
pub use dictionary_encoding::{DictionaryEncodedBytesColumn, DictionaryEncodedStrColumn};
pub use numerical::{
open_column_u64, open_column_u128, open_column_u128_as_compact_u64,
serialize_column_mappable_to_u64, serialize_column_mappable_to_u128,
};
pub(crate) use plain::{
PLAIN_BLOCK_MAX_NUM_VALUES, PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD, serialize_test_block,
PLAIN_BLOCK_MAX_NUM_VALUES, PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD, serialize_onpair_block,
};
pub use plain::{
PlainBytesColumn, PlainBytesColumnAccessor, PlainStrColumn, PlainStrColumnAccessor,
};
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,
serialize_column_mappable_to_u128,
};
pub(crate) use serialize::{open_column_bytes_from_file_slice, open_column_str_from_file_slice};
use crate::column_index::{ColumnIndex, Set};
use crate::column_values::monotonic_mapping::StrictlyMonotonicMappingToInternal;
columnar/src/column/numerical.rs
+102
View File
@@ -0,0 +1,102 @@
use std::io;
use std::io::Write;
use common::OwnedBytes;
use crate::Version;
use crate::column::Column;
use crate::column_index::{SerializableColumnIndex, serialize_column_index};
use crate::column_values::{
CodecType, MonotonicallyMappableToU64, MonotonicallyMappableToU128,
load_u64_based_column_values, serialize_column_values_u128, serialize_u64_based_column_values,
};
use crate::iterable::Iterable;
pub fn serialize_column_mappable_to_u128<T: MonotonicallyMappableToU128>(
column_index: SerializableColumnIndex<'_>,
iterable: &dyn Iterable<T>,
output: &mut impl Write,
) -> io::Result<()> {
let column_index_num_bytes = serialize_column_index(column_index, output)?;
serialize_column_values_u128(iterable, output)?;
output.write_all(&column_index_num_bytes.to_le_bytes())?;
Ok(())
}
pub fn serialize_column_mappable_to_u64<T: MonotonicallyMappableToU64>(
column_index: SerializableColumnIndex<'_>,
column_values: &impl Iterable<T>,
output: &mut impl Write,
) -> io::Result<()> {
let column_index_num_bytes = serialize_column_index(column_index, output)?;
serialize_u64_based_column_values(
column_values,
&[CodecType::Bitpacked, CodecType::BlockwiseLinear],
output,
)?;
output.write_all(&column_index_num_bytes.to_le_bytes())?;
Ok(())
}
pub fn open_column_u64<T: MonotonicallyMappableToU64>(
bytes: OwnedBytes,
format_version: Version,
) -> io::Result<Column<T>> {
let (body, column_index_num_bytes_payload) = bytes.rsplit(4);
let column_index_num_bytes = u32::from_le_bytes(
column_index_num_bytes_payload
.as_slice()
.try_into()
.unwrap(),
);
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
let column_index = crate::column_index::open_column_index(column_index_data, format_version)?;
let column_values = load_u64_based_column_values(column_values_data)?;
Ok(Column {
index: column_index,
values: column_values,
})
}
pub fn open_column_u128<T: MonotonicallyMappableToU128>(
bytes: OwnedBytes,
format_version: Version,
) -> io::Result<Column<T>> {
let (body, column_index_num_bytes_payload) = bytes.rsplit(4);
let column_index_num_bytes = u32::from_le_bytes(
column_index_num_bytes_payload
.as_slice()
.try_into()
.unwrap(),
);
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
let column_index = crate::column_index::open_column_index(column_index_data, format_version)?;
let column_values = crate::column_values::open_u128_mapped(column_values_data)?;
Ok(Column {
index: column_index,
values: column_values,
})
}
/// Open the column as u64.
///
/// See [`open_u128_as_compact_u64`] for more details.
pub fn open_column_u128_as_compact_u64(
bytes: OwnedBytes,
format_version: Version,
) -> io::Result<Column<u64>> {
let (body, column_index_num_bytes_payload) = bytes.rsplit(4);
let column_index_num_bytes = u32::from_le_bytes(
column_index_num_bytes_payload
.as_slice()
.try_into()
.unwrap(),
);
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
let column_index = crate::column_index::open_column_index(column_index_data, format_version)?;
let column_values = crate::column_values::open_u128_as_compact_u64(column_values_data)?;
Ok(Column {
index: column_index,
values: column_values,
})
}
columnar/src/column/plain.rs → columnar/src/column/plain/mod.rs
+183 -180
View File
@@ -1,4 +1,5 @@
use std::ops::Range;
use std::str::Utf8Error;
use std::sync::Arc;
use std::{fmt, io};
@@ -8,6 +9,10 @@ use onpair::{CompactDictionary, OwnedDictionaryStorage};
use crate::{Cardinality, ColumnIndex, RowId};
mod open;
pub(crate) use open::open_plain_bytes_column;
const ONPAIR_BLOCK_FOOTER_NUM_BYTES: usize = 12;
/// Target amount of uncompressed value bytes in one OnPair block.
@@ -18,15 +23,14 @@ const ONPAIR_BLOCK_FOOTER_NUM_BYTES: usize = 12;
/// boundaries to seek through. A higher threshold amortizes those dictionary, directory, and seek
/// costs over more values, but downloads more data for a point lookup and can dilute the locality
/// benefit of sorting, potentially reducing the compression ratio.
#[cfg(test)]
pub(crate) const PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD: usize = 10 * 1024 * 1024;
/// Bounds block-local offset storage when values are very short or empty.
#[cfg(test)]
pub(crate) const PLAIN_BLOCK_MAX_NUM_VALUES: usize = 262_144;
#[derive(Clone, Debug)]
pub(crate) struct PlainBlockMeta {
/// Index that helps identifying, given a value id, which block it belongs to.
#[derive(Clone)]
pub(crate) struct PlainBlockIndex {
end_bytes: Box<[usize]>,
end_values: Box<[u32]>,
}
@@ -37,7 +41,7 @@ struct PlainBlockRange {
value_range: Range<u32>,
}
impl PlainBlockMeta {
impl PlainBlockIndex {
pub(crate) fn try_new(
end_bytes: Box<[usize]>,
end_values: Box<[u32]>,
@@ -48,11 +52,24 @@ impl PlainBlockMeta {
"plain column block endpoint arrays have different lengths",
));
}
if end_bytes.last().copied().unwrap_or(0) != block_data_len {
if !end_bytes.is_sorted_by(|previous, current| previous < current) {
return Err(invalid_data(
"plain column block byte endpoints are not strictly increasing",
));
}
let previous_end_byte = end_bytes.last().copied().unwrap_or(0);
if previous_end_byte != block_data_len {
return Err(invalid_data(
"plain column block directory does not cover the block payload",
));
}
if !end_values.is_sorted_by(|previous, current| previous < current) {
return Err(invalid_data(
"plain column block value endpoints are not strictly increasing",
));
}
Ok(Self {
end_bytes,
end_values,
@@ -67,28 +84,28 @@ impl PlainBlockMeta {
self.end_values.last().copied().unwrap_or(0)
}
fn find_block(&self, value_ord: u32) -> Option<(usize, PlainBlockRange)> {
fn find_block_range(&self, value_ord: u32) -> Option<(usize, PlainBlockRange)> {
if value_ord >= self.num_values() {
return None;
}
let block_ord = self
.end_values
.partition_point(|&end_value| end_value <= value_ord);
Some((block_ord, self.block(block_ord)?))
let block: PlainBlockRange = self.block(block_ord)?;
Some((block_ord, block))
}
fn block(&self, block_ord: usize) -> Option<PlainBlockRange> {
let &end_byte = self.end_bytes.get(block_ord)?;
let &end_value = self.end_values.get(block_ord)?;
let (start_byte, start_value) = block_ord
.checked_sub(1)
.map(|previous_block_ord| {
(
self.end_bytes[previous_block_ord],
self.end_values[previous_block_ord],
)
})
.unwrap_or((0, 0));
let (start_byte, start_value) = if block_ord == 0 {
(0, 0)
} else {
(
self.end_bytes[block_ord - 1],
self.end_values[block_ord - 1],
)
};
Some(PlainBlockRange {
byte_range: start_byte..end_byte,
value_range: start_value..end_value,
@@ -98,7 +115,7 @@ impl PlainBlockMeta {
struct PlainBytesColumnData {
block_data: FileSlice,
blocks: PlainBlockMeta,
block_index: PlainBlockIndex,
}
/// A byte column whose values are stored directly rather than as dictionary ordinals.
@@ -116,7 +133,7 @@ impl fmt::Debug for PlainBytesColumn {
f.debug_struct("PlainBytesColumn")
.field("column_index", &self.column_index)
.field("num_values", &self.num_values())
.field("num_blocks", &self.data.blocks.len())
.field("num_blocks", &self.data.block_index.len())
.finish()
}
}
@@ -125,13 +142,16 @@ impl PlainBytesColumn {
pub(crate) fn open(
column_index: ColumnIndex,
block_data: FileSlice,
blocks: PlainBlockMeta,
blocks: PlainBlockIndex,
) -> io::Result<Self> {
let num_values = blocks.num_values();
validate_index_num_values(&column_index, num_values)?;
Ok(Self {
column_index,
data: Arc::new(PlainBytesColumnData { block_data, blocks }),
data: Arc::new(PlainBytesColumnData {
block_data,
block_index: blocks,
}),
})
}
@@ -139,7 +159,7 @@ impl PlainBytesColumn {
pub fn num_rows(&self) -> RowId {
match &self.column_index {
ColumnIndex::Empty { num_docs } => *num_docs,
ColumnIndex::Full => self.data.blocks.num_values(),
ColumnIndex::Full => self.data.block_index.num_values(),
ColumnIndex::Optional(optional_index) => optional_index.num_docs(),
ColumnIndex::Multivalued(multivalued_index) => multivalued_index.num_docs(),
}
@@ -147,7 +167,7 @@ impl PlainBytesColumn {
/// Returns the number of values in the column.
pub fn num_values(&self) -> u32 {
self.data.blocks.num_values()
self.data.block_index.num_values()
}
/// Returns the column index mapping rows to physical values.
@@ -173,22 +193,19 @@ impl PlainBytesColumn {
}
}
fn value_ords(&self, row_id: RowId) -> io::Result<Range<u32>> {
if row_id >= self.num_rows() {
return Err(invalid_input("plain column row id is out of bounds"));
}
Ok(self.column_index.value_row_ids(row_id))
fn value_ords(&self, row_id: RowId) -> Range<u32> {
self.column_index.value_row_ids(row_id)
}
#[cfg(test)]
pub(crate) fn for_test(column_index: ColumnIndex, values: &[&[u8]]) -> Self {
if values.is_empty() {
let blocks = PlainBlockMeta::try_new(Box::new([]), Box::new([]), 0).unwrap();
let blocks = PlainBlockIndex::try_new(Box::new([]), Box::new([]), 0).unwrap();
return Self::open(column_index, FileSlice::empty(), blocks).unwrap();
}
let block_bytes = serialize_test_block(values);
let block_len = block_bytes.len();
let blocks = PlainBlockMeta::try_new(
let blocks = PlainBlockIndex::try_new(
Box::new([block_len]),
Box::new([values.len() as u32]),
block_len,
@@ -222,37 +239,31 @@ impl fmt::Debug for PlainBytesColumnAccessor {
impl PlainBytesColumnAccessor {
/// Decodes the physical value at `value_ord` into the accessor's reusable buffer.
pub fn get_val(&mut self, value_ord: u32) -> io::Result<&[u8]> {
let Some((block_ord, block_range)) = self.column.data.blocks.find_block(value_ord) else {
return Err(invalid_input("plain column value ordinal is out of bounds"));
};
let is_cached = self
.cached_block
.as_ref()
.is_some_and(|(cached_block_ord, _)| *cached_block_ord == block_ord);
if !is_cached {
let block_bytes = self
.column
.data
.block_data
.read_bytes_slice(block_range.byte_range.clone())?;
let num_values = block_range.value_range.end - block_range.value_range.start;
let block = open_onpair_block(block_bytes, num_values)?;
self.cached_block = Some((block_ord, block));
}
pub fn get_val(&mut self, value_ord: u32) -> &[u8] {
let Self {
column,
cached_block,
output,
} = self;
let (block_ord, block_range) = column
.data
.block_index
.find_block_range(value_ord)
.expect("plain column value ordinal is out of bounds");
let block =
get_block(column, cached_block, block_ord, &block_range).expect("failed to get block");
let local_value_ord = value_ord - block_range.value_range.start;
let block = &self.cached_block.as_ref().unwrap().1;
decode_value(block, local_value_ord, &mut self.output)?;
Ok(self.output.as_slice())
decode_value(block, local_value_ord, output).expect("failed to decode value");
output.as_slice()
}
/// Decodes the first value associated with `row_id`.
pub fn first(&mut self, row_id: RowId) -> io::Result<Option<&[u8]>> {
let Some(value_ord) = self.column.value_ords(row_id)?.next() else {
return Ok(None);
pub fn first(&mut self, row_id: RowId) -> Option<&[u8]> {
let Some(value_ord) = self.column.value_ords(row_id).next() else {
return None;
};
self.get_val(value_ord).map(Some)
Some(self.get_val(value_ord))
}
/// Decodes each value associated with `row_id`, reusing the accessor's output buffer.
@@ -261,8 +272,8 @@ impl PlainBytesColumnAccessor {
row_id: RowId,
mut callback: impl FnMut(&[u8]),
) -> io::Result<()> {
for value_ord in self.column.value_ords(row_id)? {
callback(self.get_val(value_ord)?);
for value_ord in self.column.value_ords(row_id) {
callback(self.get_val(value_ord));
}
Ok(())
}
@@ -273,6 +284,29 @@ impl PlainBytesColumnAccessor {
}
}
fn get_block<'a>(
column: &PlainBytesColumn,
cached_block: &'a mut Option<(usize, onpair::Column<u32>)>,
block_ord: usize,
block_range: &PlainBlockRange,
) -> io::Result<&'a onpair::Column<u32>> {
let is_cached = if let Some((cached_block_ord, _)) = cached_block {
*cached_block_ord == block_ord
} else {
false
};
if !is_cached {
let block_bytes = column
.data
.block_data
.read_bytes_slice(block_range.byte_range.clone())?;
let num_values = block_range.value_range.end - block_range.value_range.start;
let block = open_onpair_block(block_bytes, num_values)?;
*cached_block = Some((block_ord, block));
}
Ok(&cached_block.as_ref().unwrap().1)
}
/// UTF-8 view over a [`PlainBytesColumn`].
#[derive(Clone, Debug)]
pub struct PlainStrColumn(PlainBytesColumn);
@@ -327,17 +361,15 @@ pub struct PlainStrColumnAccessor(PlainBytesColumnAccessor);
impl PlainStrColumnAccessor {
/// Decodes the physical value at `value_ord` and validates its UTF-8.
pub fn get_val(&mut self, value_ord: u32) -> io::Result<&str> {
let bytes = self.0.get_val(value_ord)?;
std::str::from_utf8(bytes).map_err(invalid_utf8)
pub fn get_val(&mut self, value_ord: u32) -> Result<&str, std::str::Utf8Error> {
let bytes = self.0.get_val(value_ord);
std::str::from_utf8(bytes)
}
/// Decodes the first value associated with `row_id` and validates its UTF-8.
pub fn first(&mut self, row_id: RowId) -> io::Result<Option<&str>> {
let Some(bytes) = self.0.first(row_id)? else {
return Ok(None);
};
std::str::from_utf8(bytes).map(Some).map_err(invalid_utf8)
pub fn first(&mut self, row_id: RowId) -> Option<&str> {
let first_bytes: &[u8] = self.0.first(row_id)?;
std::str::from_utf8(first_bytes).ok()
}
/// Decodes each value associated with `row_id`, validates its UTF-8, and invokes `callback`.
@@ -345,8 +377,8 @@ impl PlainStrColumnAccessor {
&mut self,
row_id: RowId,
mut callback: impl FnMut(&str),
) -> io::Result<()> {
for value_ord in self.0.column.value_ords(row_id)? {
) -> std::result::Result<(), Utf8Error> {
for value_ord in self.0.column.value_ords(row_id) {
callback(self.get_val(value_ord)?);
}
Ok(())
@@ -394,20 +426,15 @@ pub(crate) fn open_onpair_block(
let body_len = data.len() - ONPAIR_BLOCK_FOOTER_NUM_BYTES;
let (body, footer) = data.split(body_len);
let footer = footer.as_slice();
let dictionary_bytes_num_bytes = read_u32(&footer[0..4]) as usize;
let dictionary_offsets_num_bytes = read_u32(&footer[4..8]) as usize;
let codes_num_bytes = read_u32(&footer[8..12]) as usize;
let value_offsets_num_bytes = usize::try_from(num_values)
.ok()
.and_then(|num_values| num_values.checked_add(1))
.and_then(|num_offsets| num_offsets.checked_mul(4))
.ok_or_else(|| invalid_data("OnPair block value-offset length overflows"))?;
let dictionary_bytes_num_bytes: usize = read_u32(&footer[0..4]) as usize;
let dictionary_offsets_num_bytes: usize = read_u32(&footer[4..8]) as usize;
let codes_num_bytes: usize = read_u32(&footer[8..12]) as usize;
let value_offsets_num_bytes: usize = (1 + num_values as usize) * 4;
let expected_body_len = dictionary_bytes_num_bytes
.checked_add(dictionary_offsets_num_bytes)
.and_then(|len| len.checked_add(codes_num_bytes))
.and_then(|len| len.checked_add(value_offsets_num_bytes))
.ok_or_else(|| invalid_data("OnPair block region lengths overflow"))?;
+ dictionary_offsets_num_bytes
+ codes_num_bytes
+ value_offsets_num_bytes;
if expected_body_len != body.len() {
return Err(invalid_data(
"OnPair block region lengths do not match the block payload",
@@ -425,27 +452,27 @@ pub(crate) fn open_onpair_block(
let (dictionary_bytes, body) = body.split(dictionary_bytes_num_bytes);
let (dictionary_offsets_bytes, body) = body.split(dictionary_offsets_num_bytes);
let (codes_bytes, value_offsets_bytes) = body.split(codes_num_bytes);
let dictionary_offsets = dictionary_offsets_bytes
let dictionary_offsets: Vec<u32> = dictionary_offsets_bytes
.as_slice()
.chunks_exact(4)
.map(read_u32)
.collect::<Vec<_>>();
let codes = codes_bytes
.collect();
let codes: Vec<u16> = codes_bytes
.as_slice()
.chunks_exact(2)
.map(read_u16)
.collect::<Vec<_>>();
let value_offsets = value_offsets_bytes
.collect();
let value_offsets: Vec<u32> = value_offsets_bytes
.as_slice()
.chunks_exact(4)
.map(read_u32)
.collect::<Vec<_>>();
.collect();
let dictionary = CompactDictionary::validate(OwnedDictionaryStorage::new(
dictionary_bytes.as_slice().to_vec(),
dictionary_offsets,
))
.map_err(|error| invalid_data_owned(format!("invalid OnPair dictionary: {error}")))?;
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
let num_tokens = dictionary.num_tokens();
if codes.iter().any(|&code| code as usize >= num_tokens) {
return Err(invalid_data(
@@ -453,21 +480,13 @@ pub(crate) fn open_onpair_block(
));
}
if value_offsets.first().copied() != Some(0)
|| value_offsets
.windows(2)
.any(|offsets| offsets[0] > offsets[1])
|| value_offsets.last().copied().map(|offset| offset as usize) != Some(codes.len())
|| !value_offsets.is_sorted()
|| value_offsets.last().copied() != Some(codes.len() as u32)
{
return Err(invalid_data(
"OnPair block value offsets do not delimit the code stream",
));
}
let max_decoded_len = (codes.len() as u64)
.checked_mul(onpair::MAX_TOKEN_SIZE as u64)
.ok_or_else(|| invalid_data("OnPair block decoded length overflows"))?;
usize::try_from(max_decoded_len)
.map_err(|_| invalid_data("OnPair block decoded length does not fit in memory"))?;
Ok(onpair::Column {
dict: dictionary,
codes,
@@ -485,10 +504,12 @@ fn validate_index_num_values(column_index: &ColumnIndex, num_values: u32) -> io:
Some(start_offsets.get_val(start_offsets.num_vals() - 1))
}
};
if expected_num_values.is_some_and(|expected| expected != num_values) {
return Err(invalid_data(
"plain column index value count does not match its block directory",
));
if let Some(expected_num_values) = expected_num_values {
if expected_num_values != num_values {
return Err(invalid_data(
"plain column index value count does not match its block directory",
));
}
}
Ok(())
}
@@ -505,16 +526,40 @@ fn invalid_data(message: &'static str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message)
}
fn invalid_data_owned(message: String) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message)
}
pub(crate) fn serialize_onpair_block(
raw_bytes: &[u8],
raw_offsets: &[u32],
output: &mut impl io::Write,
) -> io::Result<u32> {
let column = onpair::compress(raw_bytes, raw_offsets, onpair::DEFAULT_CONFIG)
.map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?;
let (dictionary, codes, value_offsets) = column.into_raw();
let (dictionary_bytes, dictionary_offsets) = dictionary.into_raw();
fn invalid_input(message: &'static str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidInput, message)
}
let dictionary_offsets_num_bytes = (dictionary_offsets.len() * 4) as u32;
let codes_num_bytes = codes.len() as u32 * 2;
let dictionary_bytes_num_bytes = dictionary_bytes.len() as u32;
let value_offsets_num_bytes = value_offsets.len() as u32 * 4;
let block_num_bytes = dictionary_bytes_num_bytes
+ dictionary_offsets_num_bytes
+ codes_num_bytes
+ value_offsets_num_bytes
+ ONPAIR_BLOCK_FOOTER_NUM_BYTES as u32;
fn invalid_utf8(error: std::str::Utf8Error) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, error)
output.write_all(&dictionary_bytes)?;
for offset in dictionary_offsets {
output.write_all(&offset.to_le_bytes())?;
}
for code in codes {
output.write_all(&code.to_le_bytes())?;
}
for offset in value_offsets {
output.write_all(&offset.to_le_bytes())?;
}
output.write_all(&dictionary_bytes_num_bytes.to_le_bytes())?;
output.write_all(&dictionary_offsets_num_bytes.to_le_bytes())?;
output.write_all(&codes_num_bytes.to_le_bytes())?;
Ok(block_num_bytes)
}
#[cfg(test)]
@@ -525,24 +570,8 @@ pub(crate) fn serialize_test_block(values: &[&[u8]]) -> Vec<u8> {
raw_bytes.extend_from_slice(value);
raw_offsets.push(raw_bytes.len() as u32);
}
let column = onpair::compress(&raw_bytes, &raw_offsets, onpair::DEFAULT_CONFIG).unwrap();
let (dictionary, codes, value_offsets) = column.into_raw();
let (dictionary_bytes, dictionary_offsets) = dictionary.into_raw();
let mut output = Vec::new();
output.extend_from_slice(&dictionary_bytes);
for offset in &dictionary_offsets {
output.extend_from_slice(&offset.to_le_bytes());
}
for code in &codes {
output.extend_from_slice(&code.to_le_bytes());
}
for offset in &value_offsets {
output.extend_from_slice(&offset.to_le_bytes());
}
output.extend_from_slice(&(dictionary_bytes.len() as u32).to_le_bytes());
output.extend_from_slice(&((dictionary_offsets.len() * 4) as u32).to_le_bytes());
output.extend_from_slice(&((codes.len() * 2) as u32).to_le_bytes());
serialize_onpair_block(&raw_bytes, &raw_offsets, &mut output).unwrap();
output
}
@@ -561,22 +590,22 @@ mod tests {
#[test]
fn test_plain_block_meta_empty() {
let blocks = PlainBlockMeta::try_new(Box::new([]), Box::new([]), 0).unwrap();
let blocks = PlainBlockIndex::try_new(Box::new([]), Box::new([]), 0).unwrap();
assert_eq!(blocks.len(), 0);
assert_eq!(blocks.num_values(), 0);
assert_eq!(blocks.find_block(0), None);
assert_eq!(blocks.find_block_range(0), None);
assert_eq!(blocks.block(0), None);
}
#[test]
fn test_plain_block_meta_lookup_and_ranges() {
let blocks =
PlainBlockMeta::try_new(Box::new([10, 25, 40]), Box::new([2, 5, 9]), 40).unwrap();
PlainBlockIndex::try_new(Box::new([10, 25, 40]), Box::new([2, 5, 9]), 40).unwrap();
assert_eq!(blocks.len(), 3);
assert_eq!(blocks.num_values(), 9);
assert_eq!(
blocks.find_block(0),
blocks.find_block_range(0),
Some((
0,
PlainBlockRange {
@@ -585,9 +614,9 @@ mod tests {
}
))
);
assert_eq!(blocks.find_block(1), blocks.find_block(0));
assert_eq!(blocks.find_block_range(1), blocks.find_block_range(0));
assert_eq!(
blocks.find_block(2),
blocks.find_block_range(2),
Some((
1,
PlainBlockRange {
@@ -596,9 +625,9 @@ mod tests {
}
))
);
assert_eq!(blocks.find_block(4), blocks.find_block(2));
assert_eq!(blocks.find_block_range(4), blocks.find_block_range(2));
assert_eq!(
blocks.find_block(5),
blocks.find_block_range(5),
Some((
2,
PlainBlockRange {
@@ -607,35 +636,11 @@ mod tests {
}
))
);
assert_eq!(blocks.find_block(8), blocks.find_block(5));
assert_eq!(blocks.find_block(9), None);
assert_eq!(blocks.find_block_range(8), blocks.find_block_range(5));
assert_eq!(blocks.find_block_range(9), None);
assert_eq!(blocks.block(3), None);
}
#[test]
fn test_plain_block_meta_validation() {
let invalid_directories = [
(vec![1], vec![], 1),
(vec![0], vec![1], 0),
(vec![1], vec![0], 1),
(vec![1, 1], vec![1, 2], 1),
(vec![1, 2], vec![2, 2], 2),
(vec![1], vec![1], 2),
];
for (end_bytes, end_values, block_data_len) in invalid_directories {
assert_eq!(
PlainBlockMeta::try_new(
end_bytes.into_boxed_slice(),
end_values.into_boxed_slice(),
block_data_len,
)
.unwrap_err()
.kind(),
io::ErrorKind::InvalidData
);
}
}
#[test]
fn test_plain_bytes_full_access() {
let values: &[&[u8]] = &[b"alpha", b"", &[0, 255]];
@@ -645,18 +650,19 @@ mod tests {
assert_eq!(column.get_cardinality(), Cardinality::Full);
let mut accessor = column.accessor();
assert_eq!(accessor.get_val(0).unwrap(), b"alpha");
assert_eq!(accessor.get_val(1).unwrap(), b"");
assert_eq!(accessor.first(2).unwrap(), Some(&[0, 255][..]));
assert_eq!(accessor.get_val(0), b"alpha");
assert_eq!(accessor.get_val(1), b"");
assert_eq!(accessor.first(2), Some(&[0, 255][..]));
assert_eq!(accessor.num_cached_blocks(), 1);
assert_eq!(
accessor.get_val(3).unwrap_err().kind(),
io::ErrorKind::InvalidInput
);
assert_eq!(
accessor.first(3).unwrap_err().kind(),
io::ErrorKind::InvalidInput
);
}
#[should_panic]
#[test]
fn test_plain_bytes_panics_if_out_of_bound() {
let values: &[&[u8]] = &[b"alpha", b"", &[0, 255]];
let column = PlainBytesColumn::for_test(ColumnIndex::Full, values);
let mut accessor = column.accessor();
let _ = accessor.get_val(3);
}
#[test]
@@ -666,14 +672,14 @@ mod tests {
assert_eq!(column.num_rows(), 5);
assert_eq!(column.get_cardinality(), Cardinality::Optional);
let mut accessor = column.accessor();
assert_eq!(accessor.first(0).unwrap(), None);
assert_eq!(accessor.first(1).unwrap(), Some(&b"one"[..]));
assert_eq!(accessor.first(4).unwrap(), Some(&b"four"[..]));
assert_eq!(accessor.first(0), None);
assert_eq!(accessor.first(1), Some(&b"one"[..]));
assert_eq!(accessor.first(4), Some(&b"four"[..]));
let empty = PlainBytesColumn::for_test(ColumnIndex::Empty { num_docs: 3 }, &[]);
assert_eq!(empty.num_rows(), 3);
assert_eq!(empty.num_values(), 0);
assert_eq!(empty.accessor().first(0).unwrap(), None);
assert_eq!(empty.accessor().first(0), None);
}
#[test]
@@ -692,7 +698,7 @@ mod tests {
.for_each_value(1, |value| values.push(value.to_vec()))
.unwrap();
assert!(values.is_empty());
assert_eq!(accessor.first(2).unwrap(), Some(&b"third"[..]));
assert_eq!(accessor.first(2), Some(&b"third"[..]));
}
#[test]
@@ -704,7 +710,7 @@ mod tests {
assert_eq!(column.get_cardinality(), Cardinality::Full);
let mut accessor = column.accessor();
assert_eq!(accessor.get_val(0).unwrap(), "café");
assert_eq!(accessor.first(1).unwrap(), Some("tea"));
assert_eq!(accessor.first(1), Some("tea"));
let mut values = Vec::new();
accessor
.for_each_value(0, |value| values.push(value.to_owned()))
@@ -713,10 +719,7 @@ mod tests {
let invalid =
PlainStrColumn::wrap(PlainBytesColumn::for_test(ColumnIndex::Full, &[&[0xff]]));
assert_eq!(
invalid.accessor().get_val(0).unwrap_err().kind(),
io::ErrorKind::InvalidData
);
assert!(invalid.accessor().get_val(0).is_err());
}
#[test]
columnar/src/column/plain/open.rs
+349
View File
@@ -0,0 +1,349 @@
use std::io;
use common::HasLen;
use common::file_slice::FileSlice;
use super::{PlainBlockIndex, PlainBytesColumn};
use crate::Version;
const PLAIN_FOOTER_NUM_BYTES: usize = 8;
const PLAIN_BLOCK_INDEX_ENDPOINT_NUM_BYTES: usize = 4;
pub(crate) fn open_plain_bytes_column(data: FileSlice) -> io::Result<PlainBytesColumn> {
if data.len() < PLAIN_FOOTER_NUM_BYTES {
return Err(invalid_data("truncated plain string/byte column footer"));
}
let (body, footer_slice) = data.split_from_end(PLAIN_FOOTER_NUM_BYTES);
let footer = footer_slice.read_bytes()?;
let column_index_num_bytes = read_u32(&footer[0..4]) as usize;
let num_blocks = read_u32(&footer[4..8]) as usize;
let endpoint_array_num_bytes = num_blocks
.checked_mul(PLAIN_BLOCK_INDEX_ENDPOINT_NUM_BYTES)
.ok_or_else(|| invalid_data("plain column block directory length overflows"))?;
let directory_num_bytes = endpoint_array_num_bytes
.checked_mul(2)
.ok_or_else(|| invalid_data("plain column block directory length overflows"))?;
if directory_num_bytes > body.len() {
return Err(invalid_data(
"plain column block directory exceeds the payload",
));
}
let (column_and_blocks, directory_slice) = body.split_from_end(directory_num_bytes);
if column_index_num_bytes > column_and_blocks.len() {
return Err(invalid_data("plain column index exceeds the payload"));
}
let (column_index_slice, block_data) = column_and_blocks.split(column_index_num_bytes);
let column_index =
crate::column_index::open_column_index(column_index_slice.read_bytes()?, Version::V3)?;
let directory = directory_slice.read_bytes()?;
let (end_bytes_data, end_values_data) = directory.as_slice().split_at(endpoint_array_num_bytes);
let end_bytes: Box<[usize]> = end_bytes_data
.chunks_exact(PLAIN_BLOCK_INDEX_ENDPOINT_NUM_BYTES)
.map(|bytes| read_u32(bytes) as usize)
.collect();
let end_values: Box<[u32]> = end_values_data
.chunks_exact(PLAIN_BLOCK_INDEX_ENDPOINT_NUM_BYTES)
.map(read_u32)
.collect();
let blocks = PlainBlockIndex::try_new(end_bytes, end_values, block_data.len())?;
PlainBytesColumn::open(column_index, block_data, blocks)
}
fn read_u32(bytes: &[u8]) -> u32 {
u32::from_le_bytes(bytes.try_into().unwrap())
}
fn invalid_data(message: &'static str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message)
}
#[cfg(test)]
mod tests {
use std::ops::Range;
use std::sync::{Arc, Mutex};
use common::file_slice::{FileHandle, FileSlice};
use common::{HasLen, OwnedBytes};
use super::*;
use crate::column::{
BytesColumn, StrColumn, open_column_bytes, open_column_bytes_from_file_slice,
open_column_str,
};
use crate::column_index::{SerializableColumnIndex, serialize_column_index};
use crate::{PayloadEncoding, Version};
#[derive(Debug)]
struct RecordingFileHandle {
data: Arc<[u8]>,
reads: Mutex<Vec<Range<usize>>>,
}
impl RecordingFileHandle {
fn new(data: Vec<u8>) -> Self {
Self {
data: data.into(),
reads: Mutex::new(Vec::new()),
}
}
fn reads(&self) -> Vec<Range<usize>> {
self.reads.lock().unwrap().clone()
}
fn clear_reads(&self) {
self.reads.lock().unwrap().clear();
}
}
impl HasLen for RecordingFileHandle {
fn len(&self) -> usize {
self.data.len()
}
}
impl FileHandle for RecordingFileHandle {
fn read_bytes(&self, range: Range<usize>) -> io::Result<OwnedBytes> {
self.reads.lock().unwrap().push(range.clone());
Ok(OwnedBytes::new(self.data[range].to_vec()))
}
}
struct SerializedPlainColumn {
bytes: Vec<u8>,
block_ranges: Vec<Range<usize>>,
}
#[test]
fn test_open_v3_plain_bytes_column() {
let data = serialize_plain_column(&[b"alpha", b"", &[0, 255]]).bytes;
let BytesColumn::Plain(column) =
open_column_bytes(OwnedBytes::new(data), Version::V3).unwrap()
else {
panic!("expected a plain byte column")
};
assert_eq!(column.num_values(), 3);
let mut accessor = column.accessor();
assert_eq!(accessor.get_val(0), b"alpha");
assert_eq!(accessor.get_val(1), b"");
assert_eq!(accessor.first(2), Some(&[0, 255][..]));
}
#[test]
fn test_open_v3_plain_str_column() {
let data = serialize_plain_column(&["café".as_bytes(), b"tea"]).bytes;
let StrColumn::Plain(column) = open_column_str(OwnedBytes::new(data), Version::V3).unwrap()
else {
panic!("expected a plain string column")
};
let mut accessor = column.accessor();
assert_eq!(accessor.get_val(0).unwrap(), "café");
assert_eq!(accessor.first(1), Some("tea"));
let invalid_utf8 = serialize_plain_column(&[&[0xff]]).bytes;
let StrColumn::Plain(column) =
open_column_str(OwnedBytes::new(invalid_utf8), Version::V3).unwrap()
else {
panic!("expected a plain string column")
};
assert!(column.accessor().get_val(0).is_err());
}
#[test]
fn test_open_v3_plain_column_rejects_inconsistent_block_index() {
let mut data = serialize_plain_column(&[b"value"]).bytes;
let directory_start = data.len() - PLAIN_FOOTER_NUM_BYTES - 8;
let end_byte = read_u32(&data[directory_start..directory_start + 4]);
data[directory_start..directory_start + 4].copy_from_slice(&(end_byte + 1).to_le_bytes());
let error = open_column_bytes(OwnedBytes::new(data), Version::V3).unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn test_plain_column_block_index_serializes_endpoint_arrays() {
let serialized =
serialize_plain_column_with_limits(&[b"zero", b"one", b"two"], usize::MAX, 1);
let num_blocks = serialized.block_ranges.len();
let endpoint_array_num_bytes = num_blocks * PLAIN_BLOCK_INDEX_ENDPOINT_NUM_BYTES;
let directory_start =
serialized.bytes.len() - PLAIN_FOOTER_NUM_BYTES - endpoint_array_num_bytes * 2;
let directory =
&serialized.bytes[directory_start..directory_start + endpoint_array_num_bytes * 2];
let (end_bytes_data, end_values_data) = directory.split_at(endpoint_array_num_bytes);
let end_bytes: Vec<u32> = end_bytes_data
.chunks_exact(PLAIN_BLOCK_INDEX_ENDPOINT_NUM_BYTES)
.map(read_u32)
.collect();
let end_values: Vec<u32> = end_values_data
.chunks_exact(PLAIN_BLOCK_INDEX_ENDPOINT_NUM_BYTES)
.map(read_u32)
.collect();
let block_data_start = serialized.block_ranges[0].start;
let expected_end_bytes: Vec<u32> = serialized
.block_ranges
.iter()
.map(|block_range| (block_range.end - block_data_start) as u32)
.collect();
assert_eq!(end_bytes, expected_end_bytes);
assert_eq!(end_values, [1, 2, 3]);
}
#[test]
fn test_plain_column_loads_and_caches_only_selected_blocks() {
let values: &[&[u8]] = &[b"zero", b"one", b"two", b"three", b"four"];
let serialized = serialize_plain_column_with_limits(values, usize::MAX, 1);
assert_eq!(serialized.block_ranges.len(), 5);
let handle = Arc::new(RecordingFileHandle::new(serialized.bytes));
let BytesColumn::Plain(column) =
open_column_bytes_from_file_slice(FileSlice::new(handle.clone()), Version::V3).unwrap()
else {
panic!("expected a plain byte column")
};
let open_reads = handle.reads();
assert!(open_reads.iter().all(|read| {
serialized
.block_ranges
.iter()
.all(|block| !ranges_overlap(read, block))
}));
handle.clear_reads();
let mut accessor = column.accessor();
assert_eq!(accessor.get_val(0), b"zero");
assert_eq!(handle.reads(), [serialized.block_ranges[0].clone()]);
assert_eq!(accessor.get_val(0), b"zero");
assert_eq!(handle.reads(), [serialized.block_ranges[0].clone()]);
for (value_ord, value) in values.iter().enumerate().skip(1) {
assert_eq!(accessor.get_val(value_ord as u32), *value);
}
assert_eq!(accessor.num_cached_blocks(), 1);
assert_eq!(handle.reads().len(), 5);
assert_eq!(accessor.get_val(0), b"zero");
assert_eq!(accessor.num_cached_blocks(), 1);
assert_eq!(handle.reads().len(), 6);
assert_eq!(handle.reads().last(), Some(&serialized.block_ranges[0]));
let mut second_accessor = column.accessor();
assert_eq!(second_accessor.get_val(0), b"zero");
assert_eq!(handle.reads().len(), 7);
assert_eq!(handle.reads().last(), Some(&serialized.block_ranges[0]));
}
#[test]
fn test_plain_column_block_limits() {
assert_eq!(
super::super::PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD,
10 * 1024 * 1024
);
let values: &[&[u8]] = &[b"aa", b"bb", b"", b"cc"];
let by_bytes = serialize_plain_column_with_limits(values, 4, usize::MAX);
assert_eq!(by_bytes.block_ranges.len(), 2);
let by_values = serialize_plain_column_with_limits(values, usize::MAX, 2);
assert_eq!(by_values.block_ranges.len(), 2);
}
#[test]
fn test_open_empty_v3_plain_column() {
let data = serialize_plain_column(&[]).bytes;
let BytesColumn::Plain(column) =
open_column_bytes(OwnedBytes::new(data), Version::V3).unwrap()
else {
panic!("expected a plain byte column")
};
assert_eq!(column.num_values(), 0);
assert_eq!(column.num_rows(), 0);
}
fn ranges_overlap(left: &Range<usize>, right: &Range<usize>) -> bool {
left.start < right.end && right.start < left.end
}
fn serialize_plain_column(values: &[&[u8]]) -> SerializedPlainColumn {
serialize_plain_column_with_limits(
values,
super::super::PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD,
super::super::PLAIN_BLOCK_MAX_NUM_VALUES,
)
}
fn serialize_plain_column_with_limits(
values: &[&[u8]],
raw_num_bytes_threshold: usize,
max_num_values: usize,
) -> SerializedPlainColumn {
assert!(raw_num_bytes_threshold > 0);
assert!(max_num_values > 0);
let mut output = vec![PayloadEncoding::Plain.to_code()];
let column_index_num_bytes =
serialize_column_index(SerializableColumnIndex::Full, &mut output).unwrap();
let mut end_bytes = Vec::new();
let mut end_values = Vec::new();
let mut block_ranges = Vec::new();
let mut block_value_start = 0usize;
let mut raw_num_bytes = 0usize;
for (value_ord, value) in values.iter().enumerate() {
raw_num_bytes = raw_num_bytes.checked_add(value.len()).unwrap();
let block_num_values = value_ord + 1 - block_value_start;
if raw_num_bytes >= raw_num_bytes_threshold || block_num_values >= max_num_values {
serialize_block(
&values[block_value_start..=value_ord],
value_ord + 1,
&mut output,
&mut end_bytes,
&mut end_values,
&mut block_ranges,
);
block_value_start = value_ord + 1;
raw_num_bytes = 0;
}
}
if block_value_start < values.len() {
serialize_block(
&values[block_value_start..],
values.len(),
&mut output,
&mut end_bytes,
&mut end_values,
&mut block_ranges,
);
}
let num_blocks = block_ranges.len() as u32;
for end_byte in end_bytes {
output.extend_from_slice(&end_byte.to_le_bytes());
}
for end_value in end_values {
output.extend_from_slice(&end_value.to_le_bytes());
}
output.extend_from_slice(&column_index_num_bytes.to_le_bytes());
output.extend_from_slice(&num_blocks.to_le_bytes());
SerializedPlainColumn {
bytes: output,
block_ranges,
}
}
fn serialize_block(
values: &[&[u8]],
value_end: usize,
output: &mut Vec<u8>,
end_bytes: &mut Vec<u32>,
end_values: &mut Vec<u32>,
block_ranges: &mut Vec<Range<usize>>,
) {
let block = super::super::serialize_test_block(values);
let block_start = output.len();
output.extend_from_slice(&block);
let block_end = output.len();
block_ranges.push(block_start..block_end);
let previous_end_byte = end_bytes.last().copied().unwrap_or(0);
end_bytes.push(previous_end_byte + block.len() as u32);
end_values.push(value_end as u32);
}
}
columnar/src/column/serialize.rs
-512
View File
@@ -1,512 +0,0 @@
use std::io;
use std::io::Write;
use std::sync::Arc;
use common::file_slice::FileSlice;
use common::{HasLen, OwnedBytes};
use sstable::Dictionary;
use crate::column::{
BytesColumn, Column, DictionaryEncodedBytesColumn, DictionaryEncodedStrColumn, PlainBlockMeta,
PlainBytesColumn, PlainStrColumn,
};
use crate::column_index::{SerializableColumnIndex, serialize_column_index};
use crate::column_values::{
CodecType, MonotonicallyMappableToU64, MonotonicallyMappableToU128,
load_u64_based_column_values, serialize_column_values_u128, serialize_u64_based_column_values,
};
use crate::iterable::Iterable;
use crate::{PayloadEncoding, StrColumn, Version};
pub fn serialize_column_mappable_to_u128<T: MonotonicallyMappableToU128>(
column_index: SerializableColumnIndex<'_>,
iterable: &dyn Iterable<T>,
output: &mut impl Write,
) -> io::Result<()> {
let column_index_num_bytes = serialize_column_index(column_index, output)?;
serialize_column_values_u128(iterable, output)?;
output.write_all(&column_index_num_bytes.to_le_bytes())?;
Ok(())
}
pub fn serialize_column_mappable_to_u64<T: MonotonicallyMappableToU64>(
column_index: SerializableColumnIndex<'_>,
column_values: &impl Iterable<T>,
output: &mut impl Write,
) -> io::Result<()> {
let column_index_num_bytes = serialize_column_index(column_index, output)?;
serialize_u64_based_column_values(
column_values,
&[CodecType::Bitpacked, CodecType::BlockwiseLinear],
output,
)?;
output.write_all(&column_index_num_bytes.to_le_bytes())?;
Ok(())
}
pub fn open_column_u64<T: MonotonicallyMappableToU64>(
bytes: OwnedBytes,
format_version: Version,
) -> io::Result<Column<T>> {
let (body, column_index_num_bytes_payload) = bytes.rsplit(4);
let column_index_num_bytes = u32::from_le_bytes(
column_index_num_bytes_payload
.as_slice()
.try_into()
.unwrap(),
);
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
let column_index = crate::column_index::open_column_index(column_index_data, format_version)?;
let column_values = load_u64_based_column_values(column_values_data)?;
Ok(Column {
index: column_index,
values: column_values,
})
}
pub fn open_column_u128<T: MonotonicallyMappableToU128>(
bytes: OwnedBytes,
format_version: Version,
) -> io::Result<Column<T>> {
let (body, column_index_num_bytes_payload) = bytes.rsplit(4);
let column_index_num_bytes = u32::from_le_bytes(
column_index_num_bytes_payload
.as_slice()
.try_into()
.unwrap(),
);
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
let column_index = crate::column_index::open_column_index(column_index_data, format_version)?;
let column_values = crate::column_values::open_u128_mapped(column_values_data)?;
Ok(Column {
index: column_index,
values: column_values,
})
}
/// Open the column as u64.
///
/// See [`open_u128_as_compact_u64`] for more details.
pub fn open_column_u128_as_compact_u64(
bytes: OwnedBytes,
format_version: Version,
) -> io::Result<Column<u64>> {
let (body, column_index_num_bytes_payload) = bytes.rsplit(4);
let column_index_num_bytes = u32::from_le_bytes(
column_index_num_bytes_payload
.as_slice()
.try_into()
.unwrap(),
);
let (column_index_data, column_values_data) = body.split(column_index_num_bytes as usize);
let column_index = crate::column_index::open_column_index(column_index_data, format_version)?;
let column_values = crate::column_values::open_u128_as_compact_u64(column_values_data)?;
Ok(Column {
index: column_index,
values: column_values,
})
}
pub fn open_column_bytes(data: OwnedBytes, format_version: Version) -> io::Result<BytesColumn> {
open_column_bytes_from_file_slice(FileSlice::new(Arc::new(data)), format_version)
}
pub(crate) fn open_column_bytes_from_file_slice(
data: FileSlice,
format_version: Version,
) -> io::Result<BytesColumn> {
match format_version {
Version::V1 | Version::V2 => {
open_dictionary_bytes_column(data.read_bytes()?, format_version)
}
Version::V3 => {
if data.len() < 1 {
return Err(invalid_data("missing string/byte payload encoding tag"));
}
let (encoding_slice, payload) = data.split(1);
let encoding_bytes = encoding_slice.read_bytes()?;
let encoding =
PayloadEncoding::try_from_code(encoding_bytes[0]).map_err(io::Error::from)?;
match encoding {
PayloadEncoding::Dictionary => {
open_dictionary_bytes_column(payload.read_bytes()?, format_version)
}
PayloadEncoding::Plain => open_plain_bytes_column(payload).map(Into::into),
}
}
}
}
fn open_dictionary_bytes_column(
data: OwnedBytes,
format_version: Version,
) -> io::Result<BytesColumn> {
if data.len() < 4 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"truncated dictionary string/byte column payload",
));
}
let (body, dictionary_len_bytes) = data.rsplit(4);
let dictionary_len = u32::from_le_bytes(dictionary_len_bytes.as_slice().try_into().unwrap());
if dictionary_len as usize > body.len() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"dictionary length exceeds string/byte column payload",
));
}
let (dictionary_bytes, column_bytes) = body.split(dictionary_len as usize);
let dictionary = Arc::new(Dictionary::from_bytes(dictionary_bytes)?);
let term_ord_column = crate::column::open_column_u64::<u64>(column_bytes, format_version)?;
Ok(DictionaryEncodedBytesColumn {
dictionary,
term_ord_column,
}
.into())
}
pub fn open_column_str(data: OwnedBytes, format_version: Version) -> io::Result<StrColumn> {
open_column_str_from_file_slice(FileSlice::new(Arc::new(data)), format_version)
}
pub(crate) fn open_column_str_from_file_slice(
data: FileSlice,
format_version: Version,
) -> io::Result<StrColumn> {
match open_column_bytes_from_file_slice(data, format_version)? {
BytesColumn::DictionaryEncoded(bytes_column) => {
Ok(DictionaryEncodedStrColumn::wrap(bytes_column).into())
}
BytesColumn::Plain(bytes_column) => Ok(PlainStrColumn::wrap(bytes_column).into()),
}
}
const PLAIN_FOOTER_NUM_BYTES: usize = 8;
const PLAIN_BLOCK_DIRECTORY_ENTRY_NUM_BYTES: usize = 8;
fn open_plain_bytes_column(data: FileSlice) -> io::Result<PlainBytesColumn> {
if data.len() < PLAIN_FOOTER_NUM_BYTES {
return Err(invalid_data("truncated plain string/byte column footer"));
}
let (body, footer_slice) = data.split_from_end(PLAIN_FOOTER_NUM_BYTES);
let footer = footer_slice.read_bytes()?;
let column_index_num_bytes = read_u32(&footer[0..4]) as usize;
let num_blocks = read_u32(&footer[4..8]) as usize;
let directory_num_bytes = num_blocks
.checked_mul(PLAIN_BLOCK_DIRECTORY_ENTRY_NUM_BYTES)
.ok_or_else(|| invalid_data("plain column block directory length overflows"))?;
if directory_num_bytes > body.len() {
return Err(invalid_data(
"plain column block directory exceeds the payload",
));
}
let (column_and_blocks, directory_slice) = body.split_from_end(directory_num_bytes);
if column_index_num_bytes > column_and_blocks.len() {
return Err(invalid_data("plain column index exceeds the payload"));
}
let (column_index_slice, block_data) = column_and_blocks.split(column_index_num_bytes);
let column_index =
crate::column_index::open_column_index(column_index_slice.read_bytes()?, Version::V3)?;
let directory = directory_slice.read_bytes()?;
let mut end_bytes = Vec::with_capacity(num_blocks);
let mut end_values = Vec::with_capacity(num_blocks);
let mut byte_start = 0usize;
for entry in directory
.as_slice()
.chunks_exact(PLAIN_BLOCK_DIRECTORY_ENTRY_NUM_BYTES)
{
let block_num_bytes = read_u32(&entry[0..4]) as usize;
let value_end = read_u32(&entry[4..8]);
let byte_end = byte_start
.checked_add(block_num_bytes)
.ok_or_else(|| invalid_data("plain column block address overflows"))?;
end_bytes.push(byte_end);
end_values.push(value_end);
byte_start = byte_end;
}
let blocks = PlainBlockMeta::try_new(
end_bytes.into_boxed_slice(),
end_values.into_boxed_slice(),
block_data.len(),
)?;
PlainBytesColumn::open(column_index, block_data, blocks)
}
fn read_u32(bytes: &[u8]) -> u32 {
u32::from_le_bytes(bytes.try_into().unwrap())
}
fn invalid_data(message: &'static str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message)
}
#[cfg(test)]
mod tests {
use std::ops::Range;
use std::sync::Mutex;
use common::file_slice::FileHandle;
use super::*;
use crate::column::{
PLAIN_BLOCK_MAX_NUM_VALUES, PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD, serialize_test_block,
};
use crate::column_index::SerializableColumnIndex;
#[derive(Debug)]
struct RecordingFileHandle {
data: Arc<[u8]>,
reads: Mutex<Vec<Range<usize>>>,
}
impl RecordingFileHandle {
fn new(data: Vec<u8>) -> Self {
Self {
data: data.into(),
reads: Mutex::new(Vec::new()),
}
}
fn reads(&self) -> Vec<Range<usize>> {
self.reads.lock().unwrap().clone()
}
fn clear_reads(&self) {
self.reads.lock().unwrap().clear();
}
}
impl HasLen for RecordingFileHandle {
fn len(&self) -> usize {
self.data.len()
}
}
impl FileHandle for RecordingFileHandle {
fn read_bytes(&self, range: Range<usize>) -> io::Result<OwnedBytes> {
self.reads.lock().unwrap().push(range.clone());
Ok(OwnedBytes::new(self.data[range].to_vec()))
}
}
struct SerializedPlainColumn {
bytes: Vec<u8>,
block_ranges: Vec<Range<usize>>,
}
#[test]
fn test_v3_payload_encoding_tag_errors() {
let error = open_column_bytes(OwnedBytes::new(Vec::new()), Version::V3).unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
let error = open_column_bytes(OwnedBytes::new(vec![u8::MAX]), Version::V3).unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
let error = open_column_bytes(
OwnedBytes::new(vec![PayloadEncoding::Plain.to_code()]),
Version::V3,
)
.unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
let error = open_column_bytes(
OwnedBytes::new(vec![PayloadEncoding::Dictionary.to_code()]),
Version::V3,
)
.unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn test_open_v3_plain_bytes_column() {
let data = serialize_plain_column(&[b"alpha", b"", &[0, 255]]).bytes;
let BytesColumn::Plain(column) =
open_column_bytes(OwnedBytes::new(data), Version::V3).unwrap()
else {
panic!("expected a plain byte column")
};
assert_eq!(column.num_values(), 3);
let mut accessor = column.accessor();
assert_eq!(accessor.get_val(0).unwrap(), b"alpha");
assert_eq!(accessor.get_val(1).unwrap(), b"");
assert_eq!(accessor.first(2).unwrap(), Some(&[0, 255][..]));
}
#[test]
fn test_open_v3_plain_str_column() {
let data = serialize_plain_column(&["café".as_bytes(), b"tea"]).bytes;
let StrColumn::Plain(column) = open_column_str(OwnedBytes::new(data), Version::V3).unwrap()
else {
panic!("expected a plain string column")
};
let mut accessor = column.accessor();
assert_eq!(accessor.get_val(0).unwrap(), "café");
assert_eq!(accessor.first(1).unwrap(), Some("tea"));
let invalid_utf8 = serialize_plain_column(&[&[0xff]]).bytes;
let StrColumn::Plain(column) =
open_column_str(OwnedBytes::new(invalid_utf8), Version::V3).unwrap()
else {
panic!("expected a plain string column")
};
assert_eq!(
column.accessor().get_val(0).unwrap_err().kind(),
io::ErrorKind::InvalidData
);
}
#[test]
fn test_open_v3_plain_column_rejects_inconsistent_directory() {
let mut data = serialize_plain_column(&[b"value"]).bytes;
let directory_start = data.len() - PLAIN_FOOTER_NUM_BYTES - 8;
let block_num_bytes = read_u32(&data[directory_start..directory_start + 4]);
data[directory_start..directory_start + 4]
.copy_from_slice(&(block_num_bytes + 1).to_le_bytes());
let error = open_column_bytes(OwnedBytes::new(data), Version::V3).unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
}
#[test]
fn test_plain_column_loads_and_caches_only_selected_blocks() {
let values: &[&[u8]] = &[b"zero", b"one", b"two", b"three", b"four"];
let serialized = serialize_plain_column_with_limits(values, usize::MAX, 1);
assert_eq!(serialized.block_ranges.len(), 5);
let handle = Arc::new(RecordingFileHandle::new(serialized.bytes));
let BytesColumn::Plain(column) =
open_column_bytes_from_file_slice(FileSlice::new(handle.clone()), Version::V3).unwrap()
else {
panic!("expected a plain byte column")
};
let open_reads = handle.reads();
assert!(open_reads.iter().all(|read| {
serialized
.block_ranges
.iter()
.all(|block| !ranges_overlap(read, block))
}));
handle.clear_reads();
let mut accessor = column.accessor();
assert_eq!(accessor.get_val(0).unwrap(), b"zero");
assert_eq!(handle.reads(), [serialized.block_ranges[0].clone()]);
assert_eq!(accessor.get_val(0).unwrap(), b"zero");
assert_eq!(handle.reads(), [serialized.block_ranges[0].clone()]);
for (value_ord, value) in values.iter().enumerate().skip(1) {
assert_eq!(accessor.get_val(value_ord as u32).unwrap(), *value);
}
assert_eq!(accessor.num_cached_blocks(), 1);
assert_eq!(handle.reads().len(), 5);
assert_eq!(accessor.get_val(0).unwrap(), b"zero");
assert_eq!(accessor.num_cached_blocks(), 1);
assert_eq!(handle.reads().len(), 6);
assert_eq!(handle.reads().last(), Some(&serialized.block_ranges[0]));
let mut second_accessor = column.accessor();
assert_eq!(second_accessor.get_val(0).unwrap(), b"zero");
assert_eq!(handle.reads().len(), 7);
assert_eq!(handle.reads().last(), Some(&serialized.block_ranges[0]));
}
#[test]
fn test_plain_column_block_limits() {
assert_eq!(PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD, 10 * 1024 * 1024);
let values: &[&[u8]] = &[b"aa", b"bb", b"", b"cc"];
let by_bytes = serialize_plain_column_with_limits(values, 4, usize::MAX);
assert_eq!(by_bytes.block_ranges.len(), 2);
let by_values = serialize_plain_column_with_limits(values, usize::MAX, 2);
assert_eq!(by_values.block_ranges.len(), 2);
}
#[test]
fn test_open_empty_v3_plain_column() {
let data = serialize_plain_column(&[]).bytes;
let BytesColumn::Plain(column) =
open_column_bytes(OwnedBytes::new(data), Version::V3).unwrap()
else {
panic!("expected a plain byte column")
};
assert_eq!(column.num_values(), 0);
assert_eq!(column.num_rows(), 0);
}
fn ranges_overlap(left: &Range<usize>, right: &Range<usize>) -> bool {
left.start < right.end && right.start < left.end
}
fn serialize_plain_column(values: &[&[u8]]) -> SerializedPlainColumn {
serialize_plain_column_with_limits(
values,
PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD,
PLAIN_BLOCK_MAX_NUM_VALUES,
)
}
fn serialize_plain_column_with_limits(
values: &[&[u8]],
raw_num_bytes_threshold: usize,
max_num_values: usize,
) -> SerializedPlainColumn {
assert!(raw_num_bytes_threshold > 0);
assert!(max_num_values > 0);
let mut output = vec![PayloadEncoding::Plain.to_code()];
let column_index_num_bytes =
serialize_column_index(SerializableColumnIndex::Full, &mut output).unwrap();
let mut directory = Vec::new();
let mut block_ranges = Vec::new();
let mut block_value_start = 0usize;
let mut raw_num_bytes = 0usize;
for (value_ord, value) in values.iter().enumerate() {
raw_num_bytes = raw_num_bytes.checked_add(value.len()).unwrap();
let block_num_values = value_ord + 1 - block_value_start;
if raw_num_bytes >= raw_num_bytes_threshold || block_num_values >= max_num_values {
serialize_block(
&values[block_value_start..=value_ord],
value_ord + 1,
&mut output,
&mut directory,
&mut block_ranges,
);
block_value_start = value_ord + 1;
raw_num_bytes = 0;
}
}
if block_value_start < values.len() {
serialize_block(
&values[block_value_start..],
values.len(),
&mut output,
&mut directory,
&mut block_ranges,
);
}
let num_blocks = block_ranges.len() as u32;
output.extend_from_slice(&directory);
output.extend_from_slice(&column_index_num_bytes.to_le_bytes());
output.extend_from_slice(&num_blocks.to_le_bytes());
SerializedPlainColumn {
bytes: output,
block_ranges,
}
}
fn serialize_block(
values: &[&[u8]],
value_end: usize,
output: &mut Vec<u8>,
directory: &mut Vec<u8>,
block_ranges: &mut Vec<Range<usize>>,
) {
let block = serialize_test_block(values);
let block_start = output.len();
output.extend_from_slice(&block);
let block_end = output.len();
block_ranges.push(block_start..block_end);
directory.extend_from_slice(&(block.len() as u32).to_le_bytes());
directory.extend_from_slice(&(value_end as u32).to_le_bytes());
}
}
columnar/src/column_values/u128_based/compact_space/mod.rs
-1
View File
@@ -528,7 +528,6 @@ impl CompactSpaceDecompressor {
#[cfg(test)]
mod tests {
use itertools::Itertools;
use super::*;
columnar/src/columnar/writer/column_writers.rs
+113 -12
View File
@@ -1,10 +1,11 @@
use std::cmp::Ordering;
use std::mem::size_of;
use stacker::{ExpUnrolledLinkedList, MemoryArena};
use crate::columnar::writer::column_operation::{ColumnOperation, SymbolValue};
use crate::dictionary::{DictionaryBuilder, UnorderedId};
use crate::{Cardinality, NumericalType, NumericalValue, RowId};
use crate::dictionary::DictionaryBuilder;
use crate::{Cardinality, NumericalType, NumericalValue, PayloadEncoding, RowId};
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
#[repr(u8)]
@@ -242,8 +243,7 @@ impl NumericalColumnWriter {
#[derive(Copy, Clone)]
pub(crate) struct StrOrBytesColumnWriter {
pub(crate) dictionary_id: u32,
pub(crate) column_writer: ColumnWriter,
payload: PayloadColumnWriter,
// If true, when facing a multivalued cardinality,
// values associated to a given document will be sorted.
//
@@ -254,25 +254,112 @@ pub(crate) struct StrOrBytesColumnWriter {
pub(crate) sort_values_within_row: bool,
}
#[derive(Copy, Clone)]
pub(crate) enum PayloadColumnWriter {
Dictionary(DictionaryEncodedColumnWriter),
Plain(PlainColumnWriter),
}
#[derive(Copy, Clone)]
pub(crate) struct DictionaryEncodedColumnWriter {
pub(crate) dictionary_id: u32,
column_writer: ColumnWriter,
}
#[derive(Copy, Clone, Default)]
pub(crate) struct PlainColumnWriter {
column_writer: ColumnWriter,
pub(crate) value_store_id: u32,
}
#[derive(Default)]
pub(crate) struct PlainValueStore {
concatenated_payloads: Vec<u8>,
end_offsets: Vec<usize>,
}
impl PlainValueStore {
fn push(&mut self, value: &[u8]) -> u32 {
let value_id = self.end_offsets.len() as u32;
self.concatenated_payloads.extend_from_slice(value);
self.end_offsets.push(self.concatenated_payloads.len());
value_id
}
pub(crate) fn get(&self, value_id: u32) -> &[u8] {
let value_id = value_id as usize;
let end = self.end_offsets[value_id];
let start = if value_id == 0 {
0
} else {
self.end_offsets[value_id - 1]
};
&self.concatenated_payloads[start..end]
}
pub(crate) fn mem_usage(&self) -> usize {
self.concatenated_payloads.capacity() + self.end_offsets.capacity() * size_of::<usize>()
}
}
impl StrOrBytesColumnWriter {
pub(crate) fn with_dictionary_id(dictionary_id: u32) -> StrOrBytesColumnWriter {
pub(crate) fn dictionary(dictionary_id: u32) -> StrOrBytesColumnWriter {
StrOrBytesColumnWriter {
dictionary_id,
column_writer: Default::default(),
payload: PayloadColumnWriter::Dictionary(DictionaryEncodedColumnWriter {
dictionary_id,
column_writer: ColumnWriter::default(),
}),
sort_values_within_row: false,
}
}
pub(crate) fn plain(value_store_id: u32) -> StrOrBytesColumnWriter {
StrOrBytesColumnWriter {
payload: PayloadColumnWriter::Plain(PlainColumnWriter {
value_store_id,
..PlainColumnWriter::default()
}),
sort_values_within_row: false,
}
}
pub(crate) fn encoding(self) -> PayloadEncoding {
match self.payload {
PayloadColumnWriter::Dictionary(_) => PayloadEncoding::Dictionary,
PayloadColumnWriter::Plain(_) => PayloadEncoding::Plain,
}
}
pub(crate) fn payload(self) -> PayloadColumnWriter {
self.payload
}
pub(crate) fn column_writer(self) -> ColumnWriter {
match self.payload {
PayloadColumnWriter::Dictionary(writer) => writer.column_writer,
PayloadColumnWriter::Plain(writer) => writer.column_writer,
}
}
pub(crate) fn record_bytes(
&mut self,
doc: RowId,
bytes: &[u8],
dictionaries: &mut [DictionaryBuilder],
plain_value_stores: &mut [PlainValueStore],
arena: &mut MemoryArena,
) {
let unordered_id =
dictionaries[self.dictionary_id as usize].get_or_allocate_id(bytes, arena);
self.column_writer.record(doc, unordered_id, arena);
match &mut self.payload {
PayloadColumnWriter::Dictionary(writer) => {
let unordered_id =
dictionaries[writer.dictionary_id as usize].get_or_allocate_id(bytes, arena);
writer.column_writer.record(doc, unordered_id.0, arena);
}
PayloadColumnWriter::Plain(writer) => {
let value_id = plain_value_stores[writer.value_store_id as usize].push(bytes);
writer.column_writer.record(doc, value_id, arena);
}
}
}
pub(super) fn operation_iterator<'a>(
@@ -280,8 +367,8 @@ impl StrOrBytesColumnWriter {
arena: &MemoryArena,
old_to_new_ids: Option<&[RowId]>,
byte_buffer: &'a mut Vec<u8>,
) -> impl Iterator<Item = ColumnOperation<UnorderedId>> + 'a + use<'a> {
self.column_writer
) -> impl Iterator<Item = ColumnOperation<u32>> + 'a + use<'a> {
self.column_writer()
.operation_iterator(arena, old_to_new_ids, byte_buffer)
}
}
@@ -302,6 +389,20 @@ mod tests {
assert_eq!(delta_with_last_doc(Some(1u32), 4u32), DocumentStep::Skipped);
}
#[test]
fn test_plain_value_store() {
let mut store = PlainValueStore::default();
let first = store.push(b"same");
let empty = store.push(b"");
let duplicate = store.push(b"same");
assert_eq!((first, empty, duplicate), (0, 1, 2));
assert_eq!(store.get(first), b"same");
assert_eq!(store.get(empty), b"");
assert_eq!(store.get(duplicate), b"same");
assert!(store.mem_usage() >= 8);
}
#[track_caller]
fn test_column_writer_coercion_iter_aux(
values: impl Iterator<Item = NumericalValue>,
columnar/src/columnar/writer/mod.rs
+387 -53
View File
@@ -4,6 +4,7 @@ mod serializer;
mod value_index;
use std::io;
use std::mem::size_of;
use std::net::Ipv6Addr;
use column_operation::ColumnOperation;
@@ -13,11 +14,14 @@ use common::json_path_writer::JSON_END_OF_PATH;
pub(crate) use serializer::ColumnarSerializer;
use stacker::{Addr, ArenaHashMap, MemoryArena};
use crate::column_index::{SerializableColumnIndex, SerializableOptionalIndex};
use crate::column_index::{
SerializableColumnIndex, SerializableOptionalIndex, serialize_column_index,
};
use crate::column_values::{MonotonicallyMappableToU64, MonotonicallyMappableToU128};
use crate::columnar::column_type::ColumnType;
use crate::columnar::writer::column_writers::{
ColumnWriter, NumericalColumnWriter, StrOrBytesColumnWriter,
ColumnWriter, NumericalColumnWriter, PayloadColumnWriter, PlainValueStore,
StrOrBytesColumnWriter,
};
use crate::columnar::writer::value_index::{IndexBuilder, PreallocatedIndexBuilders};
use crate::dictionary::{DictionaryBuilder, TermIdMapping, UnorderedId};
@@ -31,6 +35,24 @@ struct SpareBuffers {
value_index_builders: PreallocatedIndexBuilders,
u64_values: Vec<u64>,
ip_addr_values: Vec<Ipv6Addr>,
plain_value_ids: Vec<u32>,
plain_block_raw_values: Vec<u8>,
plain_block_offsets: Vec<u32>,
plain_block_end_bytes: Vec<u32>,
plain_block_end_values: Vec<u32>,
}
impl SpareBuffers {
fn mem_usage(&self) -> usize {
self.value_index_builders.mem_usage()
+ self.u64_values.capacity() * size_of::<u64>()
+ self.ip_addr_values.capacity() * size_of::<Ipv6Addr>()
+ self.plain_value_ids.capacity() * size_of::<u32>()
+ self.plain_block_raw_values.capacity()
+ self.plain_block_offsets.capacity() * size_of::<u32>()
+ self.plain_block_end_bytes.capacity() * size_of::<u32>()
+ self.plain_block_end_values.capacity() * size_of::<u32>()
}
}
/// Makes it possible to create a new columnar.
@@ -57,6 +79,8 @@ pub struct ColumnarWriter {
arena: MemoryArena,
// Dictionaries used to store dictionary-encoded values.
dictionaries: Vec<DictionaryBuilder>,
// Raw stores used by plain string and byte columns.
plain_value_stores: Vec<PlainValueStore>,
buffers: SpareBuffers,
}
@@ -74,6 +98,12 @@ impl ColumnarWriter {
.iter()
.map(|dict| dict.mem_usage())
.sum::<usize>()
+ self
.plain_value_stores
.iter()
.map(PlainValueStore::mem_usage)
.sum::<usize>()
+ self.buffers.mem_usage()
}
/// Returns the list of doc ids from 0..num_docs sorted by the `sort_field`
@@ -104,18 +134,40 @@ impl ColumnarWriter {
return Vec::new();
};
let dictionary_builder = &self.dictionaries[str_or_bytes_column.dictionary_id as usize];
let term_id_mapping = dictionary_builder.build_term_id_mapping(&self.arena);
let mut symbols_buffer = Vec::new();
return collect_sort_order_from_ops(
str_or_bytes_column.operation_iterator(&self.arena, None, &mut symbols_buffer),
num_docs,
reversed,
|uid| Some(term_id_mapping.to_ord(uid).0),
None,
|a, b| a.cmp(b),
);
return match str_or_bytes_column.payload() {
PayloadColumnWriter::Dictionary(writer) => {
let dictionary_builder = &self.dictionaries[writer.dictionary_id as usize];
let term_id_mapping = dictionary_builder.build_term_id_mapping(&self.arena);
collect_sort_order_from_ops(
str_or_bytes_column.operation_iterator(
&self.arena,
None,
&mut symbols_buffer,
),
num_docs,
reversed,
|unordered_id| Some(term_id_mapping.to_ord(UnorderedId(unordered_id)).0),
None,
|a, b| a.cmp(b),
)
}
PayloadColumnWriter::Plain(writer) => {
let value_store = &self.plain_value_stores[writer.value_store_id as usize];
collect_sort_order_from_ops(
str_or_bytes_column.operation_iterator(
&self.arena,
None,
&mut symbols_buffer,
),
num_docs,
reversed,
Some,
None,
|left, right| compare_optional_plain_values(*left, *right, value_store),
)
}
};
};
let mut symbols_buffer = Vec::new();
collect_sort_order_from_ops(
@@ -156,31 +208,75 @@ impl ColumnarWriter {
column_type: ColumnType,
sort_values_within_row: bool,
) {
self.record_column_type_with_encoding(
column_name,
column_type,
sort_values_within_row,
PayloadEncoding::Dictionary,
)
.expect("column was already registered with a conflicting payload encoding");
}
/// Records a column type and the payload encoding for a string or byte column.
///
/// Dictionary encoding is the only valid encoding for other column types. Re-registering a
/// string or byte column with a different encoding returns an error.
pub fn record_column_type_with_encoding(
&mut self,
column_name: &str,
column_type: ColumnType,
sort_values_within_row: bool,
encoding: PayloadEncoding,
) -> io::Result<()> {
if sort_values_within_row {
assert!(
column_type == ColumnType::Bytes || column_type == ColumnType::Str,
"sort_values_within_row is only allowed for Bytes and Str columns",
);
match column_type {
ColumnType::Bytes | ColumnType::Str => {}
_ => {
return Err(invalid_input(
"sort_values_within_row is only allowed for Bytes and Str columns",
));
}
}
}
match column_type {
ColumnType::Str | ColumnType::Bytes => {
let (hash_map, dictionaries) = (
let (hash_map, dictionaries, plain_value_stores) = (
if column_type == ColumnType::Str {
&mut self.str_field_hash_map
} else {
&mut self.bytes_field_hash_map
},
&mut self.dictionaries,
&mut self.plain_value_stores,
);
let existing_column =
hash_map.get::<StrOrBytesColumnWriter>(column_name.as_bytes());
match existing_column {
Some(column_writer) if column_writer.encoding() != encoding => {
return Err(invalid_input(
"column was already registered with a different payload encoding",
));
}
_ => {}
}
hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<StrOrBytesColumnWriter>| {
let mut column_writer = if let Some(column_writer) = column_opt {
column_writer
} else {
let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default());
StrOrBytesColumnWriter::with_dictionary_id(dictionary_id)
match encoding {
PayloadEncoding::Dictionary => {
let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default());
StrOrBytesColumnWriter::dictionary(dictionary_id)
}
PayloadEncoding::Plain => {
let value_store_id = plain_value_stores.len() as u32;
plain_value_stores.push(PlainValueStore::default());
StrOrBytesColumnWriter::plain(value_store_id)
}
}
};
column_writer.sort_values_within_row = sort_values_within_row;
column_writer
@@ -188,18 +284,21 @@ impl ColumnarWriter {
);
}
ColumnType::Bool => {
require_dictionary_encoding(encoding)?;
self.bool_field_hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<ColumnWriter>| column_opt.unwrap_or_default(),
);
}
ColumnType::DateTime => {
require_dictionary_encoding(encoding)?;
self.datetime_field_hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<ColumnWriter>| column_opt.unwrap_or_default(),
);
}
ColumnType::I64 | ColumnType::F64 | ColumnType::U64 => {
require_dictionary_encoding(encoding)?;
let numerical_type = column_type.numerical_type().unwrap();
self.numerical_field_hash_map.mutate_or_create(
column_name.as_bytes(),
@@ -210,11 +309,15 @@ impl ColumnarWriter {
},
);
}
ColumnType::IpAddr => self.ip_addr_field_hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<ColumnWriter>| column_opt.unwrap_or_default(),
),
ColumnType::IpAddr => {
require_dictionary_encoding(encoding)?;
self.ip_addr_field_hash_map.mutate_or_create(
column_name.as_bytes(),
|column_opt: Option<ColumnWriter>| column_opt.unwrap_or_default(),
);
}
}
Ok(())
}
pub fn record_numerical<T: Into<NumericalValue> + Copy>(
@@ -275,10 +378,11 @@ impl ColumnarWriter {
}
pub fn record_str(&mut self, doc: RowId, column_name: &str, value: &str) {
let (hash_map, arena, dictionaries) = (
let (hash_map, arena, dictionaries, plain_value_stores) = (
&mut self.str_field_hash_map,
&mut self.arena,
&mut self.dictionaries,
&mut self.plain_value_stores,
);
hash_map.mutate_or_create(
column_name.as_bytes(),
@@ -287,19 +391,26 @@ impl ColumnarWriter {
// Each column has its own dictionary
let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default());
StrOrBytesColumnWriter::with_dictionary_id(dictionary_id)
StrOrBytesColumnWriter::dictionary(dictionary_id)
});
column.record_bytes(doc, value.as_bytes(), dictionaries, arena);
column.record_bytes(
doc,
value.as_bytes(),
dictionaries,
plain_value_stores,
arena,
);
column
},
);
}
pub fn record_bytes(&mut self, doc: RowId, column_name: &str, value: &[u8]) {
let (hash_map, arena, dictionaries) = (
let (hash_map, arena, dictionaries, plain_value_stores) = (
&mut self.bytes_field_hash_map,
&mut self.arena,
&mut self.dictionaries,
&mut self.plain_value_stores,
);
hash_map.mutate_or_create(
column_name.as_bytes(),
@@ -308,9 +419,9 @@ impl ColumnarWriter {
// Each column has its own dictionary
let dictionary_id = dictionaries.len() as u32;
dictionaries.push(DictionaryBuilder::default());
StrOrBytesColumnWriter::with_dictionary_id(dictionary_id)
StrOrBytesColumnWriter::dictionary(dictionary_id)
});
column.record_bytes(doc, value, dictionaries, arena);
column.record_bytes(doc, value, dictionaries, plain_value_stores, arena);
column
},
);
@@ -359,7 +470,12 @@ impl ColumnarWriter {
.map(|(column_name, addr)| (column_name, ColumnType::DateTime, addr)),
);
columns.sort_unstable_by_key(|(column_name, col_type, _)| (*column_name, *col_type));
let (arena, buffers, dictionaries) = (&self.arena, &mut self.buffers, &self.dictionaries);
let (arena, buffers, dictionaries, plain_value_stores) = (
&self.arena,
&mut self.buffers,
&self.dictionaries,
&self.plain_value_stores,
);
let mut symbol_byte_buffer: Vec<u8> = Vec::new();
for (column_name, column_type, addr) in columns {
if column_name.contains(&JSON_END_OF_PATH) {
@@ -412,27 +528,46 @@ impl ColumnarWriter {
} else {
self.str_field_hash_map.read(addr)
};
let dictionary_builder =
&dictionaries[str_or_bytes_column_writer.dictionary_id as usize];
let cardinality = str_or_bytes_column_writer
.column_writer
.column_writer()
.get_cardinality(num_docs);
let mut column_serializer =
serializer.start_serialize_column(column_name, column_type);
serialize_bytes_or_str_column(
cardinality,
num_docs,
str_or_bytes_column_writer.sort_values_within_row,
dictionary_builder,
str_or_bytes_column_writer.operation_iterator(
arena,
old_to_new_row_ids,
&mut symbol_byte_buffer,
),
buffers,
&self.arena,
&mut column_serializer,
)?;
match str_or_bytes_column_writer.payload() {
PayloadColumnWriter::Dictionary(writer) => {
let dictionary_builder = &dictionaries[writer.dictionary_id as usize];
serialize_dictionary_bytes_or_str_column(
cardinality,
num_docs,
str_or_bytes_column_writer.sort_values_within_row,
dictionary_builder,
str_or_bytes_column_writer.operation_iterator(
arena,
old_to_new_row_ids,
&mut symbol_byte_buffer,
),
buffers,
arena,
&mut column_serializer,
)?;
}
PayloadColumnWriter::Plain(writer) => {
let value_store = &plain_value_stores[writer.value_store_id as usize];
serialize_plain_bytes_or_str_column(
cardinality,
num_docs,
str_or_bytes_column_writer.sort_values_within_row,
value_store,
str_or_bytes_column_writer.operation_iterator(
arena,
old_to_new_row_ids,
&mut symbol_byte_buffer,
),
buffers,
&mut column_serializer,
)?;
}
}
column_serializer.finalize()?;
}
ColumnType::F64 | ColumnType::I64 | ColumnType::U64 => {
@@ -535,12 +670,12 @@ fn collect_sort_order_from_ops<V, K: Clone>(
// V3 serialize [PayloadEncoding, Dictionary, Column, dictionary num bytes U32::LE]
// Column: [Column Index, Column Values, column index num bytes U32::LE]
#[expect(clippy::too_many_arguments)]
fn serialize_bytes_or_str_column(
fn serialize_dictionary_bytes_or_str_column(
cardinality: Cardinality,
num_docs: RowId,
sort_values_within_row: bool,
dictionary_builder: &DictionaryBuilder,
operation_it: impl Iterator<Item = ColumnOperation<UnorderedId>>,
operation_it: impl Iterator<Item = ColumnOperation<u32>>,
buffers: &mut SpareBuffers,
arena: &MemoryArena,
wrt: impl io::Write,
@@ -557,11 +692,11 @@ fn serialize_bytes_or_str_column(
dictionary_builder.serialize(arena, &mut counting_writer)?;
let dictionary_num_bytes: u32 = counting_writer.written_bytes() as u32;
let mut wrt = counting_writer.finish();
let operation_iterator = operation_it.map(|symbol: ColumnOperation<UnorderedId>| {
let operation_iterator = operation_it.map(|symbol: ColumnOperation<u32>| {
// We map unordered ids to ordered ids.
match symbol {
ColumnOperation::Value(unordered_id) => {
let ordered_id = term_id_mapping.to_ord(unordered_id);
let ordered_id = term_id_mapping.to_ord(UnorderedId(unordered_id));
ColumnOperation::Value(ordered_id.0 as u64)
}
ColumnOperation::NewDoc(doc) => ColumnOperation::NewDoc(doc),
@@ -580,6 +715,192 @@ fn serialize_bytes_or_str_column(
Ok(())
}
fn serialize_plain_bytes_or_str_column(
cardinality: Cardinality,
num_docs: RowId,
sort_values_within_row: bool,
value_store: &PlainValueStore,
operation_it: impl Iterator<Item = ColumnOperation<u32>>,
buffers: &mut SpareBuffers,
mut wrt: impl io::Write,
) -> io::Result<()> {
let SpareBuffers {
value_index_builders,
plain_value_ids,
plain_block_raw_values,
plain_block_offsets,
plain_block_end_bytes,
plain_block_end_values,
..
} = buffers;
plain_value_ids.clear();
wrt.write_all(&[PayloadEncoding::Plain.to_code()])?;
let serializable_column_index = match cardinality {
Cardinality::Full => {
consume_operation_iterator(
operation_it,
value_index_builders.borrow_required_index_builder(),
plain_value_ids,
);
SerializableColumnIndex::Full
}
Cardinality::Optional => {
let optional_index_builder = value_index_builders.borrow_optional_index_builder();
consume_operation_iterator(operation_it, optional_index_builder, plain_value_ids);
let optional_index = optional_index_builder.finish(num_docs);
SerializableColumnIndex::Optional(SerializableOptionalIndex {
non_null_row_ids: Box::new(optional_index),
num_rows: num_docs,
})
}
Cardinality::Multivalued => {
let multivalued_index_builder = value_index_builders.borrow_multivalued_index_builder();
consume_operation_iterator(operation_it, multivalued_index_builder, plain_value_ids);
let serializable_multivalued_index = multivalued_index_builder.finish(num_docs);
if sort_values_within_row {
sort_plain_values_within_row(
serializable_multivalued_index.start_offsets.boxed_iter(),
plain_value_ids,
value_store,
);
}
SerializableColumnIndex::Multivalued(serializable_multivalued_index)
}
};
let column_index_num_bytes = serialize_column_index(serializable_column_index, &mut wrt)?;
serialize_plain_blocks(
plain_value_ids,
value_store,
plain_block_raw_values,
plain_block_offsets,
plain_block_end_bytes,
plain_block_end_values,
&mut wrt,
)?;
debug_assert_eq!(plain_block_end_bytes.len(), plain_block_end_values.len());
// Keep the serialized layout identical to PlainBlockIndex: all byte endpoints first,
// followed by all value endpoints.
for &end_byte in plain_block_end_bytes.iter() {
wrt.write_all(&end_byte.to_le_bytes())?;
}
for &end_value in plain_block_end_values.iter() {
wrt.write_all(&end_value.to_le_bytes())?;
}
let num_blocks = u32::try_from(plain_block_end_bytes.len())
.map_err(|_| invalid_input("plain column contains more than u32::MAX blocks"))?;
wrt.write_all(&column_index_num_bytes.to_le_bytes())?;
wrt.write_all(&num_blocks.to_le_bytes())?;
Ok(())
}
fn serialize_plain_blocks(
value_ids: &[u32],
value_store: &PlainValueStore,
block_raw_values: &mut Vec<u8>,
block_offsets: &mut Vec<u32>,
block_end_bytes: &mut Vec<u32>,
block_end_values: &mut Vec<u32>,
output: &mut impl io::Write,
) -> io::Result<()> {
block_raw_values.clear();
block_offsets.clear();
block_offsets.push(0);
block_end_bytes.clear();
block_end_values.clear();
let mut end_byte = 0u32;
let mut end_value = 0u32;
for &value_id in value_ids {
let value = value_store.get(value_id);
block_raw_values.extend_from_slice(value);
let block_end = u32::try_from(block_raw_values.len())
.map_err(|_| invalid_input("plain value exceeds the OnPair block size limit"))?;
block_offsets.push(block_end);
let block_num_values = block_offsets.len() - 1;
if block_raw_values.len() >= crate::column::PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD
|| block_num_values >= crate::column::PLAIN_BLOCK_MAX_NUM_VALUES
{
flush_plain_block(
block_raw_values,
block_offsets,
&mut end_byte,
&mut end_value,
block_end_bytes,
block_end_values,
output,
)?;
}
}
if block_offsets.len() > 1 {
flush_plain_block(
block_raw_values,
block_offsets,
&mut end_byte,
&mut end_value,
block_end_bytes,
block_end_values,
output,
)?;
}
Ok(())
}
fn flush_plain_block(
block_raw_values: &mut Vec<u8>,
block_offsets: &mut Vec<u32>,
end_byte: &mut u32,
end_value: &mut u32,
block_end_bytes: &mut Vec<u32>,
block_end_values: &mut Vec<u32>,
output: &mut impl io::Write,
) -> io::Result<()> {
let block_num_values = u32::try_from(block_offsets.len() - 1)
.map_err(|_| invalid_input("plain OnPair block contains too many values"))?;
let block_num_bytes =
crate::column::serialize_onpair_block(block_raw_values, block_offsets, output)?;
*end_byte = end_byte
.checked_add(block_num_bytes)
.ok_or_else(|| invalid_input("plain column block data exceeds u32::MAX bytes"))?;
*end_value = end_value
.checked_add(block_num_values)
.ok_or_else(|| invalid_input("plain column contains more than u32::MAX values"))?;
block_end_bytes.push(*end_byte);
block_end_values.push(*end_value);
block_raw_values.clear();
block_offsets.clear();
block_offsets.push(0);
Ok(())
}
fn sort_plain_values_within_row(
multivalued_index: impl Iterator<Item = RowId>,
values: &mut [u32],
value_store: &PlainValueStore,
) {
let mut start_index = 0usize;
for end_index in multivalued_index {
let end_index = end_index as usize;
values[start_index..end_index]
.sort_unstable_by(|left, right| value_store.get(*left).cmp(value_store.get(*right)));
start_index = end_index;
}
}
fn compare_optional_plain_values(
left: Option<u32>,
right: Option<u32>,
value_store: &PlainValueStore,
) -> std::cmp::Ordering {
match (left, right) {
(None, None) => std::cmp::Ordering::Equal,
(None, Some(_)) => std::cmp::Ordering::Less,
(Some(_), None) => std::cmp::Ordering::Greater,
(Some(left), Some(right)) => value_store.get(left).cmp(value_store.get(right)),
}
}
fn serialize_numerical_column(
cardinality: Cardinality,
num_docs: RowId,
@@ -817,6 +1138,19 @@ fn consume_operation_iterator<T: Ord, TIndexBuilder: IndexBuilder>(
}
}
fn require_dictionary_encoding(encoding: PayloadEncoding) -> io::Result<()> {
if encoding == PayloadEncoding::Dictionary {
return Ok(());
}
Err(invalid_input(
"plain payload encoding is only supported for Bytes and Str columns",
))
}
fn invalid_input(message: &'static str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidInput, message)
}
#[cfg(test)]
mod tests {
use stacker::MemoryArena;
columnar/src/columnar/writer/value_index.rs
+8
View File
@@ -1,3 +1,5 @@
use std::mem::size_of;
use crate::RowId;
use crate::column_index::{SerializableMultivalueIndex, SerializableOptionalIndex};
use crate::iterable::Iterable;
@@ -117,6 +119,12 @@ pub struct PreallocatedIndexBuilders {
}
impl PreallocatedIndexBuilders {
pub fn mem_usage(&self) -> usize {
self.optional_index_builder.docs.capacity() * size_of::<RowId>()
+ self.multivalued_index_builder.doc_with_values.capacity() * size_of::<RowId>()
+ self.multivalued_index_builder.start_offsets.capacity() * size_of::<u32>()
}
pub fn borrow_required_index_builder(&mut self) -> &mut FullIndexBuilder {
&mut self.required_index_builder
}
columnar/src/tests.rs
+192
View File
@@ -282,6 +282,198 @@ fn test_dictionary_encoded_bytes() {
assert_eq!(term_buffer, b"b");
}
#[test]
fn test_plain_str_roundtrip_with_row_mapping() {
let mut buffer = Vec::new();
let mut columnar_writer = ColumnarWriter::default();
columnar_writer
.record_column_type_with_encoding(
"plain",
ColumnType::Str,
false,
crate::PayloadEncoding::Plain,
)
.unwrap();
columnar_writer.record_str(0, "plain", "zero");
columnar_writer.record_str(1, "plain", "one");
columnar_writer.record_str(2, "plain", "two");
columnar_writer
.serialize(3, Some(&[2, 0, 1]), &mut buffer)
.unwrap();
let columnar_reader = ColumnarReader::open(buffer).unwrap();
let columns = columnar_reader.read_columns("plain").unwrap();
let DynamicColumn::Str(column) = columns[0].open().unwrap() else {
panic!();
};
assert_eq!(column.payload_encoding(), crate::PayloadEncoding::Plain);
let column = column.as_plain().unwrap();
assert_eq!(column.num_rows(), 3);
assert_eq!(column.get_cardinality(), Cardinality::Full);
let mut accessor = column.accessor();
assert_eq!(accessor.first(0), Some("one"));
assert_eq!(accessor.first(1), Some("two"));
assert_eq!(accessor.first(2), Some("zero"));
}
#[test]
fn test_plain_bytes_optional_and_non_utf8_roundtrip() {
let mut buffer = Vec::new();
let mut columnar_writer = ColumnarWriter::default();
columnar_writer
.record_column_type_with_encoding(
"plain",
ColumnType::Bytes,
false,
crate::PayloadEncoding::Plain,
)
.unwrap();
columnar_writer.record_bytes(1, "plain", b"");
columnar_writer.record_bytes(3, "plain", &[0, 255]);
columnar_writer.serialize(5, None, &mut buffer).unwrap();
let columnar_reader = ColumnarReader::open(buffer).unwrap();
let columns = columnar_reader.read_columns("plain").unwrap();
let DynamicColumn::Bytes(column) = columns[0].open().unwrap() else {
panic!();
};
let column = column.as_plain().unwrap();
assert_eq!(column.get_cardinality(), Cardinality::Optional);
let mut accessor = column.accessor();
assert_eq!(accessor.first(0), None);
assert_eq!(accessor.first(1), Some(&b""[..]));
assert_eq!(accessor.first(2), None);
assert_eq!(accessor.first(3), Some(&[0, 255][..]));
assert_eq!(accessor.first(4), None);
}
#[test]
fn test_plain_bytes_multivalued_values_are_sorted_within_row() {
let mut buffer = Vec::new();
let mut columnar_writer = ColumnarWriter::default();
columnar_writer
.record_column_type_with_encoding(
"plain",
ColumnType::Bytes,
true,
crate::PayloadEncoding::Plain,
)
.unwrap();
columnar_writer.record_bytes(0, "plain", b"z");
columnar_writer.record_bytes(0, "plain", b"");
columnar_writer.record_bytes(0, "plain", b"a");
columnar_writer.record_bytes(2, "plain", b"last");
columnar_writer.serialize(3, None, &mut buffer).unwrap();
let columnar_reader = ColumnarReader::open(buffer).unwrap();
let columns = columnar_reader.read_columns("plain").unwrap();
let DynamicColumn::Bytes(column) = columns[0].open().unwrap() else {
panic!();
};
let column = column.as_plain().unwrap();
assert_eq!(column.get_cardinality(), Cardinality::Multivalued);
let mut accessor = column.accessor();
let mut values = Vec::new();
accessor
.for_each_value(0, |value| values.push(value.to_vec()))
.unwrap();
assert_eq!(values, [b"".to_vec(), b"a".to_vec(), b"z".to_vec()]);
values.clear();
accessor
.for_each_value(1, |value| values.push(value.to_vec()))
.unwrap();
assert!(values.is_empty());
assert_eq!(accessor.first(2), Some(&b"last"[..]));
}
#[test]
fn test_plain_empty_column_roundtrip() {
let mut buffer = Vec::new();
let mut columnar_writer = ColumnarWriter::default();
columnar_writer
.record_column_type_with_encoding(
"plain",
ColumnType::Str,
false,
crate::PayloadEncoding::Plain,
)
.unwrap();
columnar_writer.serialize(3, None, &mut buffer).unwrap();
let columnar_reader = ColumnarReader::open(buffer).unwrap();
let columns = columnar_reader.read_columns("plain").unwrap();
let DynamicColumn::Str(column) = columns[0].open().unwrap() else {
panic!();
};
let column = column.as_plain().unwrap();
assert_eq!(column.num_values(), 0);
assert_eq!(column.accessor().first(0), None);
}
#[test]
fn test_plain_column_registration_rejects_conflicting_encoding() {
let mut columnar_writer = ColumnarWriter::default();
columnar_writer
.record_column_type_with_encoding(
"plain",
ColumnType::Str,
false,
crate::PayloadEncoding::Plain,
)
.unwrap();
columnar_writer
.record_column_type_with_encoding(
"plain",
ColumnType::Str,
true,
crate::PayloadEncoding::Plain,
)
.unwrap();
assert_eq!(
columnar_writer
.record_column_type_with_encoding(
"plain",
ColumnType::Str,
false,
crate::PayloadEncoding::Dictionary,
)
.unwrap_err()
.kind(),
std::io::ErrorKind::InvalidInput
);
assert_eq!(
columnar_writer
.record_column_type_with_encoding(
"number",
ColumnType::U64,
false,
crate::PayloadEncoding::Plain,
)
.unwrap_err()
.kind(),
std::io::ErrorKind::InvalidInput
);
}
#[test]
fn test_sort_order_plain_str() {
let mut columnar_writer = ColumnarWriter::default();
columnar_writer
.record_column_type_with_encoding(
"plain",
ColumnType::Str,
false,
crate::PayloadEncoding::Plain,
)
.unwrap();
columnar_writer.record_str(0, "plain", "z");
columnar_writer.record_str(2, "plain", "a");
columnar_writer.record_str(3, "plain", "m");
assert_eq!(columnar_writer.sort_order("plain", 4, false), [1, 2, 3, 0]);
assert_eq!(columnar_writer.sort_order("plain", 4, true), [0, 3, 2, 1]);
}
#[test]
fn test_sort_order_str_asc_desc() {
let mut dataframe_writer = ColumnarWriter::default();
plain_string_column_plan.md
+86 -51
View File
@@ -170,29 +170,23 @@ The equivalent methods should exist on `StrColumn`.
### Plain value access
OnPair16 decompression needs an output buffer, so use caller-owned scratch space and return a
slice borrowed from it:
OnPair16 decompression needs an output buffer, so create a mutable accessor that owns both its
decode buffer and its most recently parsed block:
```rust
fn get_val<'a>(
&self,
value_ord: u32,
output: &'a mut Vec<u8>,
) -> io::Result<&'a [u8]>;
let mut accessor = column.accessor();
fn first<'a>(
&self,
row_id: RowId,
output: &'a mut Vec<u8>,
) -> io::Result<Option<&'a [u8]>>;
accessor.get_val(value_ord) -> io::Result<&[u8]>;
accessor.first(row_id) -> io::Result<Option<&[u8]>>;
```
`PlainStrColumn` exposes the corresponding methods returning `&str`.
`PlainStrColumnAccessor` exposes the corresponding methods returning `&str`. Returned values are
valid until the next mutable call on that accessor.
A normal `Iterator<Item = &[u8]>` cannot safely reuse a single mutable decompression buffer for a
multivalued row. Offer either:
- a callback-based `for_each_value(row_id, scratch, callback)` API, or
- a callback-based `for_each_value(row_id, callback)` API, or
- an iterator over physical value ordinals plus `get_val()`.
The callback API is preferable for the common case because it does not expose ordinals as a
@@ -203,26 +197,31 @@ logical part of the plain-column API.
`PlainBytesColumn` contains:
- A `ColumnIndex` mapping document rows to physical value positions.
- The OnPair16 decoder/model shared by the column.
- Concatenated independently compressed payloads.
- A monotonic offsets column containing `num_values + 1` entries.
- A range-readable `FileSlice` containing independently trained OnPair blocks.
- A resident block directory storing cumulative byte and value endpoints in separate boxed slices.
Each mutable accessor owns its most recently parsed, still-compressed block and reusable output
buffer. The underlying column remains immutable and shareable.
To read value `i`:
1. Read offsets `i` and `i + 1`.
2. Slice that range from the compressed payload.
3. Decode it into the caller's scratch buffer.
4. Return `&[u8]`, or validate and return `&str` through `PlainStrColumn`.
1. Binary-search the directory's contiguous cumulative `end_value` slice.
2. Range-read the selected block unless it is the cached block.
3. Read the block-local offsets for `i` and slice its native `u16` OnPair codes.
4. Decode only that value into the accessor's output buffer.
5. Return `&[u8]`, or validate and return `&str` through `PlainStrColumn`.
The values must be compressed independently so accessing one row does not require decoding an
entire block or neighboring values. This matches OnPair16's random-access model.
Fetching a value downloads its containing block but does not decompress the entire block or any
neighboring value. Column clones share the immutable block directory, while accessors have
independent caches.
### Validation
Opening or reading a plain column must reject:
- Unknown encoding discriminants.
- Truncated model or payload regions.
- Truncated directories, blocks, or native OnPair buffers.
- Non-contiguous block byte ranges or non-increasing cumulative value ordinals.
- Non-monotonic or out-of-range offsets.
- An offsets count inconsistent with the column index/value count.
- Invalid OnPair16 tokens or model references.
@@ -266,24 +265,52 @@ layout is:
```text
1u8
| column_index
| onpair16_model
| compressed_values
| value_offsets
| onpair_block_0 ... onpair_block_n
| block_directory[]
| column_index_num_bytes:u32 LE
| model_num_bytes:u32 LE
| compressed_values_num_bytes:u64 LE
| value_offsets_num_bytes:u32 LE
| num_values:u32 LE
| num_blocks:u32 LE
```
`value_offsets` is a serialized monotonic `u64` column with exactly `num_values + 1` entries. Its
first entry is zero, its last entry equals `compressed_values_num_bytes`, and every adjacent pair
delimits one independently compressed value. `onpair16_model` uses the codec's canonical model
serialization, which must be frozen alongside the standalone plain-column implementation.
Each eight-byte block-directory entry is:
The fixed 24-byte footer is read from the end first. The four regions are then split from left to
right without copying. Checked conversion to `usize`, checked length sums, and region bounds are
required before any `OwnedBytes::split` call; unrecognized tags and trailing bytes are invalid.
```text
block_num_bytes:u32 LE | end_value:u32 LE
```
`block_num_bytes` is the serialized size of one independently loadable block. Block addresses are
derived by checked cumulative addition, so the entire column may exceed `u32::MAX` bytes even
though an individual block may not. `end_value` is the exclusive cumulative physical value
ordinal. The final entry yields the column's value count, so it is not serialized separately. An
empty column has no blocks.
Each OnPair block is:
```text
dictionary_bytes_with_read_padding
| dictionary_offsets:u32 LE[]
| codes:u16 LE[]
| value_offsets:u32 LE[]
| dictionary_bytes_num_bytes:u32 LE
| dictionary_offsets_num_bytes:u32 LE
| codes_num_bytes:u32 LE
```
The directory determines the block's value count, so `value_offsets` must contain exactly one more
entry than the block has values. Its entries index the native code stream and its final entry equals
the number of codes. Native buffers are copied into aligned typed vectors after the block is fetched
because a `FileSlice` result does not guarantee `u16` or `u32` alignment.
Opening reads the fixed eight-byte footer and directory from the end, then eagerly opens only the
global column index. Block region lengths, cumulative addresses, native-buffer alignment, OnPair
invariants, codes, and local offsets are validated before decoding.
Blocks close after appending a complete value when they reach either
`PLAIN_BLOCK_RAW_NUM_BYTES_THRESHOLD` (10 MiB of uncompressed value bytes) or
`PLAIN_BLOCK_MAX_NUM_VALUES` (262,144 values). Lowering the byte threshold reduces point-read
downloads and can improve compression when sorting makes blocks locally homogeneous, but produces
more dictionaries, directory entries, and block seeks. Raising it amortizes those costs over more
values, while increasing point-read downloads and potentially diluting sorting locality and its
compression benefit. The value-count cap independently bounds offsets for short or empty values.
V1 and V2 never consume a tag. V3 always consumes exactly one tag byte for string and byte
payloads, including dictionary payloads.
@@ -331,23 +358,25 @@ encodings.
### Plain writer flow
1. Store raw input values in the memory arena rather than interning them in the term dictionary.
2. Continue recording new-document/value operations so the existing cardinality and column-index
builders can be reused.
1. Store raw input values in a dedicated contiguous per-column store rather than interning them in
the term dictionary. Keep cumulative offsets beside the bytes so sorting and block serialization
can borrow values directly without copying the entire column out of the memory arena.
2. Continue recording new-document/value operations in the memory arena so the existing
cardinality and column-index builders can be reused.
3. Apply `old_to_new_row_ids` before serialization.
4. Sort values lexicographically within a row when requested.
5. Train the OnPair16 model from the field values or the codec's prescribed sample.
6. Compress each value independently in physical value order.
7. Append each compressed value and record the next offset.
8. Serialize the index, model, payload, offsets, footer, and V3 encoding tag.
5. Accumulate complete physical values until the 10 MiB raw-byte threshold or 262,144-value cap is
reached.
6. Train and encode an independent native OnPair column for the block.
7. Serialize the block and append its byte length and cumulative `end_value` to the directory.
8. Serialize the index, blocks, directory, footer, and V3 encoding tag.
Include raw value storage, OnPair training structures, compressed buffers, and offsets in
`mem_usage()`.
There is currently no OnPair dependency in the workspace. Implementation therefore requires
either adding the intended Rust codec dependency or introducing a small internal codec module. The
column implementation should depend on a narrow train/serialize/open/compress/decompress interface
so the storage code is not coupled to training internals.
Tantivy and `tantivy-columnar` use Rust 1.91 and depend on OnPair 0.2. The reader preserves OnPair's
native dictionary, `u16` code, and `u32` block-local offset representation and uses its validation
and random-access decoder.
## 6. Tantivy writer integration
@@ -361,8 +390,8 @@ During `FastFieldsWriter::from_schema_and_tokenizer_manager`:
Tokenizer behavior is independent of payload encoding: tokenization happens first, and each
resulting token is then recorded using the selected encoding.
Index-time sorting also runs before serialization. The plain writer should compare raw arena bytes
directly, while the dictionary writer retains the current ordered-term-ID optimization.
Index-time sorting also runs before serialization. The plain writer should compare raw value-store
bytes directly, while the dictionary writer retains the current ordered-term-ID optimization.
## 7. Columnar and segment merging
@@ -529,6 +558,12 @@ cost separately from value compression.
10. Adapt ordinal-dependent queries, collectors, sorting, and aggregations according to that
audit.
11. Complete corruption tests, compatibility-index tests, and benchmarks.
12. Improve cold OnPair block-opening performance. Profile the cost of range acquisition,
native-buffer materialization, and validation separately; avoid eagerly copying and validating
the entire code and value-offset streams for a point lookup; investigate retaining the block as
`OwnedBytes`, lazily reading only the selected value's offsets and codes, borrowing dictionary
bytes, and adding an alignment-independent zero-copy serialized view to OnPair if partial
zero-copy is insufficient.
## Existing worktree note