From 082074040a68bb1fd2f1c129464b699d3fbcd4c1 Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Tue, 11 Aug 2026 10:28:27 +0200 Subject: [PATCH] Added support for writing plain strings --- Cargo.toml | 2 +- ... dictionary_encoded_str_search_and_get.rs} | 0 columnar/src/column/bytes_or_str.rs | 84 ++- .../mod.rs} | 6 +- .../src/column/dictionary_encoding/open.rs | 35 ++ columnar/src/column/mod.rs | 23 +- columnar/src/column/numerical.rs | 102 ++++ .../src/column/{plain.rs => plain/mod.rs} | 363 +++++++------ columnar/src/column/plain/open.rs | 349 ++++++++++++ columnar/src/column/serialize.rs | 512 ------------------ .../u128_based/compact_space/mod.rs | 1 - .../src/columnar/writer/column_writers.rs | 125 ++++- columnar/src/columnar/writer/mod.rs | 440 +++++++++++++-- columnar/src/columnar/writer/value_index.rs | 8 + columnar/src/tests.rs | 192 +++++++ plain_string_column_plan.md | 137 +++-- 16 files changed, 1553 insertions(+), 826 deletions(-) rename benches/{str_search_and_get.rs => dictionary_encoded_str_search_and_get.rs} (100%) rename columnar/src/column/{dictionary_encoded.rs => dictionary_encoding/mod.rs} (97%) create mode 100644 columnar/src/column/dictionary_encoding/open.rs create mode 100644 columnar/src/column/numerical.rs rename columnar/src/column/{plain.rs => plain/mod.rs} (69%) create mode 100644 columnar/src/column/plain/open.rs delete mode 100644 columnar/src/column/serialize.rs diff --git a/Cargo.toml b/Cargo.toml index b0aea7bd4..229b2bec3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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]] diff --git a/benches/str_search_and_get.rs b/benches/dictionary_encoded_str_search_and_get.rs similarity index 100% rename from benches/str_search_and_get.rs rename to benches/dictionary_encoded_str_search_and_get.rs diff --git a/columnar/src/column/bytes_or_str.rs b/columnar/src/column/bytes_or_str.rs index 727fac489..5f0c8bcf0 100644 --- a/columnar/src/column/bytes_or_str.rs +++ b/columnar/src/column/bytes_or_str.rs @@ -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 { + 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 { + 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 { + 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 { + 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 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(); diff --git a/columnar/src/column/dictionary_encoded.rs b/columnar/src/column/dictionary_encoding/mod.rs similarity index 97% rename from columnar/src/column/dictionary_encoded.rs rename to columnar/src/column/dictionary_encoding/mod.rs index 8132fe713..6e1d38633 100644 --- a/columnar/src/column/dictionary_encoded.rs +++ b/columnar/src/column/dictionary_encoding/mod.rs @@ -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() diff --git a/columnar/src/column/dictionary_encoding/open.rs b/columnar/src/column/dictionary_encoding/open.rs new file mode 100644 index 000000000..07f5ec37b --- /dev/null +++ b/columnar/src/column/dictionary_encoding/open.rs @@ -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 { + 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::(column_bytes, format_version)?; + Ok(DictionaryEncodedBytesColumn { + dictionary, + term_ord_column, + }) +} diff --git a/columnar/src/column/mod.rs b/columnar/src/column/mod.rs index a68892004..782594495 100644 --- a/columnar/src/column/mod.rs +++ b/columnar/src/column/mod.rs @@ -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; diff --git a/columnar/src/column/numerical.rs b/columnar/src/column/numerical.rs new file mode 100644 index 000000000..f39bd8960 --- /dev/null +++ b/columnar/src/column/numerical.rs @@ -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( + column_index: SerializableColumnIndex<'_>, + iterable: &dyn Iterable, + 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( + column_index: SerializableColumnIndex<'_>, + column_values: &impl Iterable, + 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( + bytes: OwnedBytes, + format_version: Version, +) -> io::Result> { + 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( + bytes: OwnedBytes, + format_version: Version, +) -> io::Result> { + 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> { + 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, + }) +} diff --git a/columnar/src/column/plain.rs b/columnar/src/column/plain/mod.rs similarity index 69% rename from columnar/src/column/plain.rs rename to columnar/src/column/plain/mod.rs index b6ed3e7e7..6ae372f85 100644 --- a/columnar/src/column/plain.rs +++ b/columnar/src/column/plain/mod.rs @@ -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, } -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 { 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 { 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> { - 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 { + 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> { - 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)>, + block_ord: usize, + block_range: &PlainBlockRange, +) -> io::Result<&'a onpair::Column> { + 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> { - 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 = dictionary_offsets_bytes .as_slice() .chunks_exact(4) .map(read_u32) - .collect::>(); - let codes = codes_bytes + .collect(); + let codes: Vec = codes_bytes .as_slice() .chunks_exact(2) .map(read_u16) - .collect::>(); - let value_offsets = value_offsets_bytes + .collect(); + let value_offsets: Vec = value_offsets_bytes .as_slice() .chunks_exact(4) .map(read_u32) - .collect::>(); + .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 { + 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 { 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] diff --git a/columnar/src/column/plain/open.rs b/columnar/src/column/plain/open.rs new file mode 100644 index 000000000..93c648493 --- /dev/null +++ b/columnar/src/column/plain/open.rs @@ -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 { + 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>>, + } + + impl RecordingFileHandle { + fn new(data: Vec) -> Self { + Self { + data: data.into(), + reads: Mutex::new(Vec::new()), + } + } + + fn reads(&self) -> Vec> { + 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) -> io::Result { + self.reads.lock().unwrap().push(range.clone()); + Ok(OwnedBytes::new(self.data[range].to_vec())) + } + } + + struct SerializedPlainColumn { + bytes: Vec, + block_ranges: Vec>, + } + + #[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 = end_bytes_data + .chunks_exact(PLAIN_BLOCK_INDEX_ENDPOINT_NUM_BYTES) + .map(read_u32) + .collect(); + let end_values: Vec = 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 = 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, right: &Range) -> 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, + end_bytes: &mut Vec, + end_values: &mut Vec, + block_ranges: &mut Vec>, + ) { + 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); + } +} diff --git a/columnar/src/column/serialize.rs b/columnar/src/column/serialize.rs deleted file mode 100644 index afc2cf648..000000000 --- a/columnar/src/column/serialize.rs +++ /dev/null @@ -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( - column_index: SerializableColumnIndex<'_>, - iterable: &dyn Iterable, - 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( - column_index: SerializableColumnIndex<'_>, - column_values: &impl Iterable, - 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( - bytes: OwnedBytes, - format_version: Version, -) -> io::Result> { - 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( - bytes: OwnedBytes, - format_version: Version, -) -> io::Result> { - 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> { - 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 { - 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 { - 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 { - 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::(column_bytes, format_version)?; - Ok(DictionaryEncodedBytesColumn { - dictionary, - term_ord_column, - } - .into()) -} - -pub fn open_column_str(data: OwnedBytes, format_version: Version) -> io::Result { - 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 { - 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 { - 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>>, - } - - impl RecordingFileHandle { - fn new(data: Vec) -> Self { - Self { - data: data.into(), - reads: Mutex::new(Vec::new()), - } - } - - fn reads(&self) -> Vec> { - 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) -> io::Result { - self.reads.lock().unwrap().push(range.clone()); - Ok(OwnedBytes::new(self.data[range].to_vec())) - } - } - - struct SerializedPlainColumn { - bytes: Vec, - block_ranges: Vec>, - } - - #[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, right: &Range) -> 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, - directory: &mut Vec, - block_ranges: &mut Vec>, - ) { - 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()); - } -} diff --git a/columnar/src/column_values/u128_based/compact_space/mod.rs b/columnar/src/column_values/u128_based/compact_space/mod.rs index 9fa36e020..a1eda5d59 100644 --- a/columnar/src/column_values/u128_based/compact_space/mod.rs +++ b/columnar/src/column_values/u128_based/compact_space/mod.rs @@ -528,7 +528,6 @@ impl CompactSpaceDecompressor { #[cfg(test)] mod tests { - use itertools::Itertools; use super::*; diff --git a/columnar/src/columnar/writer/column_writers.rs b/columnar/src/columnar/writer/column_writers.rs index 6d3e2a415..ff9308a7c 100644 --- a/columnar/src/columnar/writer/column_writers.rs +++ b/columnar/src/columnar/writer/column_writers.rs @@ -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, + end_offsets: Vec, +} + +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::() + } +} + 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, - ) -> impl Iterator> + 'a + use<'a> { - self.column_writer + ) -> impl Iterator> + '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, diff --git a/columnar/src/columnar/writer/mod.rs b/columnar/src/columnar/writer/mod.rs index abe8d15ca..af32e980a 100644 --- a/columnar/src/columnar/writer/mod.rs +++ b/columnar/src/columnar/writer/mod.rs @@ -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, ip_addr_values: Vec, + plain_value_ids: Vec, + plain_block_raw_values: Vec, + plain_block_offsets: Vec, + plain_block_end_bytes: Vec, + plain_block_end_values: Vec, +} + +impl SpareBuffers { + fn mem_usage(&self) -> usize { + self.value_index_builders.mem_usage() + + self.u64_values.capacity() * size_of::() + + self.ip_addr_values.capacity() * size_of::() + + self.plain_value_ids.capacity() * size_of::() + + self.plain_block_raw_values.capacity() + + self.plain_block_offsets.capacity() * size_of::() + + self.plain_block_end_bytes.capacity() * size_of::() + + self.plain_block_end_values.capacity() * size_of::() + } } /// 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, + // Raw stores used by plain string and byte columns. + plain_value_stores: Vec, buffers: SpareBuffers, } @@ -74,6 +98,12 @@ impl ColumnarWriter { .iter() .map(|dict| dict.mem_usage()) .sum::() + + self + .plain_value_stores + .iter() + .map(PlainValueStore::mem_usage) + .sum::() + + 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::(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| { 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| 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| 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| 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| column_opt.unwrap_or_default(), + ); + } } + Ok(()) } pub fn record_numerical + 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 = 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( // 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>, + operation_it: impl Iterator>, 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| { + let operation_iterator = operation_it.map(|symbol: ColumnOperation| { // 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>, + 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, + block_offsets: &mut Vec, + block_end_bytes: &mut Vec, + block_end_values: &mut Vec, + 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, + block_offsets: &mut Vec, + end_byte: &mut u32, + end_value: &mut u32, + block_end_bytes: &mut Vec, + block_end_values: &mut Vec, + 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, + 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, + right: Option, + 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( } } +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; diff --git a/columnar/src/columnar/writer/value_index.rs b/columnar/src/columnar/writer/value_index.rs index d289346b8..45c28ba46 100644 --- a/columnar/src/columnar/writer/value_index.rs +++ b/columnar/src/columnar/writer/value_index.rs @@ -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::() + + self.multivalued_index_builder.doc_with_values.capacity() * size_of::() + + self.multivalued_index_builder.start_offsets.capacity() * size_of::() + } + pub fn borrow_required_index_builder(&mut self) -> &mut FullIndexBuilder { &mut self.required_index_builder } diff --git a/columnar/src/tests.rs b/columnar/src/tests.rs index 83e45612b..2a34724a4 100644 --- a/columnar/src/tests.rs +++ b/columnar/src/tests.rs @@ -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(); diff --git a/plain_string_column_plan.md b/plain_string_column_plan.md index 92cc155bf..69e918700 100644 --- a/plain_string_column_plan.md +++ b/plain_string_column_plan.md @@ -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, -) -> io::Result<&'a [u8]>; +let mut accessor = column.accessor(); -fn first<'a>( - &self, - row_id: RowId, - output: &'a mut Vec, -) -> io::Result>; +accessor.get_val(value_ord) -> io::Result<&[u8]>; +accessor.first(row_id) -> io::Result>; ``` -`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` 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