This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-mosaic.git
The following commit(s) were added to refs/heads/main by this push:
new b36d1a9 perf(core): read non-nested columns without creating Arrow
arrays (#76)
b36d1a9 is described below
commit b36d1a9300d8dabe17db22e28690f88c92d65a82
Author: jianguotian <[email protected]>
AuthorDate: Thu Aug 20 15:06:03 2026 +0800
perf(core): read non-nested columns without creating Arrow arrays (#76)
---
core/src/bucket_reader.rs | 557 ++++++++++++++++++++++++++++++++++++++++++++++
core/src/reader.rs | 79 ++++++-
core/src/reader_tests.rs | 527 +++++++++++++++++++++++++++++++++++++++++++
3 files changed, 1161 insertions(+), 2 deletions(-)
diff --git a/core/src/bucket_reader.rs b/core/src/bucket_reader.rs
index 4dbaab4..8ea8bde 100644
--- a/core/src/bucket_reader.rs
+++ b/core/src/bucket_reader.rs
@@ -23,11 +23,335 @@ use arrow_array::*;
use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer,
ScalarBuffer};
use arrow_schema::{DataType, Field, TimeUnit};
+use crate::reader::Encoding;
use crate::spec::*;
use crate::types;
use crate::values::Value;
use crate::varint;
+/// Borrowed view of one encoded scalar column.
+///
+/// The view borrows buffers owned by a [`crate::reader::RowGroupReader`] and
must be consumed
+/// synchronously during
[`crate::reader::RowGroupReader::visit_encoded_columns`]. It exposes the
+/// physical encoding without first materializing an Arrow array.
+#[derive(Clone, Copy)]
+pub struct EncodedColumn<'a> {
+ data_type: &'a DataType,
+ encoding: u8,
+ has_nulls: bool,
+ null_bitmap: &'a [u8],
+ const_value: &'a Value,
+ dict_values: &'a [Value],
+ dict_bit_width: usize,
+ data: &'a [u8],
+ data_cursor: usize,
+ num_rows: usize,
+}
+
+impl<'a> EncodedColumn<'a> {
+ #[allow(clippy::too_many_arguments)]
+ fn new(
+ data_type: &'a DataType,
+ encoding: u8,
+ has_nulls: bool,
+ null_bitmap: &'a [u8],
+ const_value: &'a Value,
+ dict_values: &'a [Value],
+ dict_bit_width: usize,
+ data: &'a [u8],
+ data_cursor: usize,
+ num_rows: usize,
+ ) -> Self {
+ Self {
+ data_type,
+ encoding,
+ has_nulls,
+ null_bitmap,
+ const_value,
+ dict_values,
+ dict_bit_width,
+ data,
+ data_cursor,
+ num_rows,
+ }
+ }
+
+ pub fn data_type(&self) -> &DataType {
+ self.data_type
+ }
+
+ /// Returns the physical encoding used by this column.
+ pub fn encoding(&self) -> Encoding {
+ Encoding::from_code(self.encoding)
+ }
+
+ /// Returns the logical row count.
+ pub fn num_rows(&self) -> usize {
+ self.num_rows
+ }
+
+ /// Returns whether at least one row is null.
+ pub fn has_nulls(&self) -> bool {
+ self.has_nulls || self.encoding == ENCODING_ALL_NULL
+ }
+
+ /// Returns the physical null bitmap when one is present.
+ ///
+ /// A set bit means that the corresponding row is null.
+ pub fn null_bitmap(&self) -> Option<&'a [u8]> {
+ self.has_nulls.then_some(self.null_bitmap)
+ }
+
+ /// Returns whether `row` is null.
+ ///
+ /// # Panics
+ ///
+ /// Panics if `row >= self.num_rows()`.
+ pub fn is_null(&self, row: usize) -> bool {
+ assert!(row < self.num_rows, "encoded column row out of bounds");
+ self.encoding == ENCODING_ALL_NULL || (self.has_nulls &&
is_null(self.null_bitmap, row))
+ }
+
+ /// Returns the single encoded value for a CONST column.
+ pub fn constant(&self) -> io::Result<Option<EncodedValueRef<'a>>> {
+ if self.encoding != ENCODING_CONST {
+ return Ok(None);
+ }
+ encoded_value_ref(self.data_type, self.const_value).map(Some)
+ }
+
+ /// Iterates values in logical row order without materializing an Arrow
array.
+ ///
+ /// Null rows are returned as [`EncodedValueRef::Null`]. Dictionary
indexes and PLAIN values
+ /// are decoded lazily as the iterator advances.
+ pub fn values(&self) -> EncodedColumnValues<'a> {
+ EncodedColumnValues {
+ column: *self,
+ row: 0,
+ data_cursor: self.data_cursor,
+ bit_offset: 0,
+ done: false,
+ }
+ }
+}
+
+/// Borrowed scalar value returned by [`EncodedColumnValues`].
+#[derive(Clone, Copy, Debug, PartialEq)]
+#[non_exhaustive]
+pub enum EncodedValueRef<'a> {
+ Null,
+ Boolean(bool),
+ Int8(i8),
+ Int16(i16),
+ Int32(i32),
+ Int64(i64),
+ Float32(f32),
+ Float64(f64),
+ Utf8(&'a [u8]),
+ Binary(&'a [u8]),
+ DecimalCompact(i64),
+ DecimalLarge(&'a [u8]),
+ Date32(i32),
+ Time32(i32),
+ TimestampMillis(i64),
+ TimestampMicros(i64),
+ TimestampNanos { millis: i64, nanos_of_milli: i32 },
+}
+
+/// Iterator over an [`EncodedColumn`] in logical row order.
+pub struct EncodedColumnValues<'a> {
+ column: EncodedColumn<'a>,
+ row: usize,
+ data_cursor: usize,
+ bit_offset: usize,
+ done: bool,
+}
+
+impl<'a> Iterator for EncodedColumnValues<'a> {
+ type Item = io::Result<EncodedValueRef<'a>>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ if self.done || self.row >= self.column.num_rows {
+ return None;
+ }
+
+ let row = self.row;
+ self.row += 1;
+ if self.column.is_null(row) {
+ return Some(Ok(EncodedValueRef::Null));
+ }
+
+ let value = match self.column.encoding {
+ ENCODING_CONST => encoded_value_ref(self.column.data_type,
self.column.const_value),
+ ENCODING_DICT => read_bit_packed_checked(
+ self.column.data,
+ self.column.data_cursor,
+ self.bit_offset,
+ self.column.dict_bit_width,
+ )
+ .and_then(|index| {
+ self.bit_offset += self.column.dict_bit_width;
+ self.column
+ .dict_values
+ .get(index)
+ .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData,
"corrupt dict index"))
+ .and_then(|value| encoded_value_ref(self.column.data_type,
value))
+ }),
+ ENCODING_PLAIN => {
+ let result =
+ read_encoded_value(self.column.data_type,
self.column.data, self.data_cursor);
+ if let Ok((_, size)) = result {
+ self.data_cursor += size;
+ }
+ result.map(|(value, _)| value)
+ }
+ ENCODING_ALL_NULL => Ok(EncodedValueRef::Null),
+ _ => Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!("unsupported encoding {}", self.column.encoding),
+ )),
+ };
+ if value.is_err() {
+ self.done = true;
+ }
+ Some(value)
+ }
+}
+
+fn encoded_value_ref<'a>(
+ data_type: &DataType,
+ value: &'a Value,
+) -> io::Result<EncodedValueRef<'a>> {
+ let value = match value {
+ Value::Null => EncodedValueRef::Null,
+ Value::Boolean(value) => EncodedValueRef::Boolean(*value),
+ Value::TinyInt(value) => EncodedValueRef::Int8(*value),
+ Value::SmallInt(value) => EncodedValueRef::Int16(*value),
+ Value::Integer(value) => EncodedValueRef::Int32(*value),
+ Value::BigInt(value) => EncodedValueRef::Int64(*value),
+ Value::Float(value) => EncodedValueRef::Float32(*value),
+ Value::Double(value) => EncodedValueRef::Float64(*value),
+ Value::Date(value) => EncodedValueRef::Date32(*value),
+ Value::Time(value) => EncodedValueRef::Time32(*value),
+ Value::String(value) => EncodedValueRef::Utf8(value),
+ Value::Bytes(value) => EncodedValueRef::Binary(value),
+ Value::DecimalCompact(value) =>
EncodedValueRef::DecimalCompact(*value),
+ Value::DecimalLarge(value) => EncodedValueRef::DecimalLarge(value),
+ Value::TimestampMillis(value) =>
EncodedValueRef::TimestampMillis(*value),
+ Value::TimestampMicros(value) =>
EncodedValueRef::TimestampMicros(*value),
+ Value::TimestampNanos {
+ millis,
+ nanos_of_milli,
+ } => EncodedValueRef::TimestampNanos {
+ millis: *millis,
+ nanos_of_milli: *nanos_of_milli,
+ },
+ };
+ validate_encoded_value(data_type, value)
+}
+
+fn validate_encoded_value<'a>(
+ data_type: &DataType,
+ value: EncodedValueRef<'a>,
+) -> io::Result<EncodedValueRef<'a>> {
+ if let EncodedValueRef::TimestampNanos {
+ millis,
+ nanos_of_milli,
+ } = value
+ {
+ if types::is_timestamp_nanos(data_type) {
+ types::millis_nanos_to_ns(millis, nanos_of_milli)?;
+ }
+ }
+ Ok(value)
+}
+
+fn read_encoded_value<'a>(
+ data_type: &DataType,
+ data: &'a [u8],
+ position: usize,
+) -> io::Result<(EncodedValueRef<'a>, usize)> {
+ let width = types::fixed_width(data_type);
+ if width <= 0 {
+ let mut payload = position;
+ let length = varint::decode(data, &mut payload).map_err(|_| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ "truncated varint in variable-length value",
+ )
+ })? as usize;
+ let end = payload
+ .checked_add(length)
+ .filter(|end| *end <= data.len())
+ .ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ "buffer truncated in variable-length value",
+ )
+ })?;
+ let value = match data_type {
+ DataType::Utf8 => EncodedValueRef::Utf8(&data[payload..end]),
+ DataType::Binary => EncodedValueRef::Binary(&data[payload..end]),
+ DataType::Decimal128(_, _) =>
EncodedValueRef::DecimalLarge(&data[payload..end]),
+ _ => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!("unsupported encoded data type: {data_type:?}"),
+ ));
+ }
+ };
+ return Ok((value, end - position));
+ }
+
+ let width = width as usize;
+ let end = position
+ .checked_add(width)
+ .filter(|end| *end <= data.len())
+ .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "column data
truncated"))?;
+ let bytes = &data[position..end];
+ let value = match data_type {
+ DataType::Boolean => EncodedValueRef::Boolean(bytes[0] != 0),
+ DataType::Int8 => EncodedValueRef::Int8(bytes[0] as i8),
+ DataType::Int16 =>
EncodedValueRef::Int16(i16::from_be_bytes([bytes[0], bytes[1]])),
+ DataType::Int32 => {
+ EncodedValueRef::Int32(i32::from_be_bytes([bytes[0], bytes[1],
bytes[2], bytes[3]]))
+ }
+ DataType::Date32 => {
+ EncodedValueRef::Date32(i32::from_be_bytes([bytes[0], bytes[1],
bytes[2], bytes[3]]))
+ }
+ DataType::Time32(_) => {
+ EncodedValueRef::Time32(i32::from_be_bytes([bytes[0], bytes[1],
bytes[2], bytes[3]]))
+ }
+ DataType::Float32 =>
EncodedValueRef::Float32(f32::from_bits(u32::from_be_bytes([
+ bytes[0], bytes[1], bytes[2], bytes[3],
+ ]))),
+ DataType::Int64 => EncodedValueRef::Int64(read_i64(data, position)),
+ DataType::Float64 =>
EncodedValueRef::Float64(f64::from_bits(read_u64(data, position))),
+ DataType::Decimal128(_, _) =>
EncodedValueRef::DecimalCompact(read_i64(data, position)),
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ EncodedValueRef::TimestampMillis(read_i64(data, position))
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ EncodedValueRef::TimestampMicros(read_i64(data, position))
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) | DataType::Struct(_)
+ if types::is_timestamp_nanos(data_type) =>
+ {
+ EncodedValueRef::TimestampNanos {
+ millis: read_i64(data, position),
+ nanos_of_milli: i32::from_be_bytes([bytes[8], bytes[9],
bytes[10], bytes[11]]),
+ }
+ }
+ _ => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!("unsupported encoded data type: {data_type:?}"),
+ ));
+ }
+ };
+ validate_encoded_value(data_type, value).map(|value| (value, width))
+}
+
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DataVariant {
Boolean,
@@ -917,6 +1241,30 @@ pub struct BucketReader {
}
impl BucketReader {
+ pub(crate) fn encoded_column(&self, column: usize) ->
io::Result<EncodedColumn<'_>> {
+ if column >= self.total_columns {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ format!(
+ "column index {} out of range (num_columns={})",
+ column, self.total_columns
+ ),
+ ));
+ }
+ Ok(EncodedColumn::new(
+ &self.col_types[column],
+ self.encodings[column],
+ self.has_nulls[column],
+ &self.null_bitmaps[column],
+ &self.const_values[column],
+ &self.dict_values[column],
+ self.dict_bit_widths[column],
+ &self.data,
+ self.data_cursors[column],
+ self.col_num_rows(column),
+ ))
+ }
+
fn col_num_rows(&self, col: usize) -> usize {
if col < self.num_primary {
self.num_rows
@@ -1201,6 +1549,21 @@ pub struct ColumnPageReader {
}
impl ColumnPageReader {
+ pub(crate) fn encoded_column(&self) -> EncodedColumn<'_> {
+ EncodedColumn::new(
+ &self.col_type,
+ self.encoding,
+ self.has_nulls,
+ &self.null_bitmap,
+ &self.const_value,
+ &self.dict_values,
+ self.dict_bit_width,
+ &self.data,
+ self.data_cursor,
+ self.num_rows,
+ )
+ }
+
pub fn new(
col_type: DataType,
encoding: u8,
@@ -1229,6 +1592,15 @@ impl ColumnPageReader {
page_data_start: usize,
num_rows: usize,
) -> io::Result<Self> {
+ if !matches!(
+ encoding,
+ ENCODING_PLAIN | ENCODING_CONST | ENCODING_DICT | ENCODING_ALL_NULL
+ ) {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!("column page: unsupported encoding {}", encoding),
+ ));
+ }
if page_data_start > data.len() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
@@ -1996,6 +2368,23 @@ fn read_bit_packed(buf: &[u8], byte_base: usize,
bit_offset: usize, bit_width: u
value
}
+fn read_bit_packed_checked(
+ buf: &[u8],
+ byte_base: usize,
+ bit_offset: usize,
+ bit_width: usize,
+) -> io::Result<usize> {
+ let bit_end = bit_offset
+ .checked_add(bit_width)
+ .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "dict index
offset overflow"))?;
+ let byte_end = byte_base
+ .checked_add(bit_end.div_ceil(8))
+ .filter(|end| *end <= buf.len())
+ .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "truncated
dict indexes"))?;
+ let _ = byte_end;
+ Ok(read_bit_packed(buf, byte_base, bit_offset, bit_width))
+}
+
fn bit_width(num_entries: usize) -> usize {
if num_entries <= 1 {
return 0;
@@ -2416,3 +2805,171 @@ mod tests {
.all(|value| value.to_bits() == 0.0f32.to_bits()));
}
}
+
+#[cfg(test)]
+mod encoded_column_tests {
+ use super::*;
+
+ #[test]
+ fn column_page_rejects_unknown_encoding_without_panicking() {
+ let result = std::panic::catch_unwind(|| -> io::Result<()> {
+ let page =
+ ColumnPageReader::new(DataType::Int32, 0xff, true,
Value::Null, Vec::new(), 1)?;
+ page.read_all()?;
+ Ok(())
+ });
+
+ let err = result
+ .expect("unknown column page encoding must not panic")
+ .expect_err("unknown column page encoding must be rejected");
+ assert_eq!(err.kind(), io::ErrorKind::InvalidData);
+ assert!(err.to_string().contains("unsupported encoding 255"));
+ }
+
+ #[test]
+ fn rejects_invalid_timestamp_nanos_for_all_encodings() {
+ let data_type = DataType::Timestamp(TimeUnit::Nanosecond, None);
+ let invalid = Value::TimestampNanos {
+ millis: 0,
+ nanos_of_milli: 1_000_000,
+ };
+ let placeholder = Value::Null;
+
+ let constant = EncodedColumn::new(
+ &data_type,
+ ENCODING_CONST,
+ false,
+ &[],
+ &invalid,
+ &[],
+ 0,
+ &[],
+ 0,
+ 1,
+ );
+ assert_eq!(
+ constant.constant().unwrap_err().kind(),
+ io::ErrorKind::InvalidData
+ );
+ assert_eq!(
+ constant.values().next().unwrap().unwrap_err().kind(),
+ io::ErrorKind::InvalidData
+ );
+
+ let dictionary = EncodedColumn::new(
+ &data_type,
+ ENCODING_DICT,
+ false,
+ &[],
+ &placeholder,
+ std::slice::from_ref(&invalid),
+ 0,
+ &[],
+ 0,
+ 1,
+ );
+ assert_eq!(
+ dictionary.values().next().unwrap().unwrap_err().kind(),
+ io::ErrorKind::InvalidData
+ );
+
+ let mut plain_data = 0i64.to_be_bytes().to_vec();
+ plain_data.extend_from_slice(&1_000_000i32.to_be_bytes());
+ let plain = EncodedColumn::new(
+ &data_type,
+ ENCODING_PLAIN,
+ false,
+ &[],
+ &placeholder,
+ &[],
+ 0,
+ &plain_data,
+ 0,
+ 1,
+ );
+ assert_eq!(
+ plain.values().next().unwrap().unwrap_err().kind(),
+ io::ErrorKind::InvalidData
+ );
+ }
+
+ #[test]
+ fn reports_corrupt_dict_indexes_without_panicking() {
+ let data_type = DataType::Int32;
+ let placeholder = Value::Null;
+ let dict_values = [Value::Integer(7)];
+
+ let truncated = EncodedColumn::new(
+ &data_type,
+ ENCODING_DICT,
+ false,
+ &[],
+ &placeholder,
+ &dict_values,
+ 1,
+ &[],
+ 0,
+ 1,
+ );
+ assert_eq!(
+ truncated.values().next().unwrap().unwrap_err().kind(),
+ io::ErrorKind::InvalidData
+ );
+
+ let out_of_range = EncodedColumn::new(
+ &data_type,
+ ENCODING_DICT,
+ false,
+ &[],
+ &placeholder,
+ &dict_values,
+ 1,
+ &[1],
+ 0,
+ 1,
+ );
+ assert_eq!(
+ out_of_range.values().next().unwrap().unwrap_err().kind(),
+ io::ErrorKind::InvalidData
+ );
+ }
+
+ #[test]
+ fn reports_truncated_plain_values_without_panicking() {
+ let placeholder = Value::Null;
+
+ let fixed = EncodedColumn::new(
+ &DataType::Int64,
+ ENCODING_PLAIN,
+ false,
+ &[],
+ &placeholder,
+ &[],
+ 0,
+ &[0; 7],
+ 0,
+ 1,
+ );
+ assert_eq!(
+ fixed.values().next().unwrap().unwrap_err().kind(),
+ io::ErrorKind::InvalidData
+ );
+
+ let variable = EncodedColumn::new(
+ &DataType::Utf8,
+ ENCODING_PLAIN,
+ false,
+ &[],
+ &placeholder,
+ &[],
+ 0,
+ &[3, b'a', b'b'],
+ 0,
+ 1,
+ );
+ assert_eq!(
+ variable.values().next().unwrap().unwrap_err().kind(),
+ io::ErrorKind::InvalidData
+ );
+ }
+}
diff --git a/core/src/reader.rs b/core/src/reader.rs
index 54ffa9d..4de3fee 100644
--- a/core/src/reader.rs
+++ b/core/src/reader.rs
@@ -22,6 +22,7 @@ use arrow_array::{ArrayRef, RecordBatch, RecordBatchOptions};
use arrow_schema::{DataType, Field, Schema};
use crate::bucket_reader::{read_typed_value, read_variable_value,
BucketReader, ColumnPageReader};
+pub use crate::bucket_reader::{EncodedColumn, EncodedColumnValues,
EncodedValueRef};
use crate::schema::MosaicSchema;
use crate::spec::*;
use crate::stats::{self, ColumnStats};
@@ -224,8 +225,8 @@ pub enum Encoding {
}
impl Encoding {
- fn from_code(c: u8) -> Self {
- match c {
+ pub(crate) fn from_code(code: u8) -> Self {
+ match code {
ENCODING_PLAIN => Encoding::Plain,
ENCODING_CONST => Encoding::Const,
ENCODING_DICT => Encoding::Dict,
@@ -1503,6 +1504,80 @@ impl RowGroupReader {
self.num_rows
}
+ /// Visits projected scalar columns in output order without materializing
Arrow arrays.
+ ///
+ /// Each [`EncodedColumn`] borrows this row group's buffers and is valid
only for the duration
+ /// of the callback. ARRAY and MAP columns are rejected before the first
callback because they
+ /// are represented by multiple physical columns.
+ pub fn visit_encoded_columns<F>(&self, mut visitor: F) -> io::Result<()>
+ where
+ F: FnMut(&str, &DataType, bool, EncodedColumn<'_>) -> io::Result<()>,
+ {
+ for &global_index in &self.output_order {
+ if self.projected_columns[global_index]
+ && matches!(
+ self.schema.columns[global_index].data_type,
+ DataType::List(_) | DataType::Map(_, _)
+ )
+ {
+ return Err(io::Error::new(
+ io::ErrorKind::Unsupported,
+ format!(
+ "encoded access does not support nested column '{}'",
+ self.schema.columns[global_index].name
+ ),
+ ));
+ }
+ }
+
+ let mut visited = vec![false; self.num_columns];
+ for &global_index in &self.output_order {
+ if visited[global_index] {
+ continue;
+ }
+ visited[global_index] = true;
+ if !self.projected_columns[global_index] {
+ continue;
+ }
+
+ let column = &self.schema.columns[global_index];
+ let bucket_id = column.bucket_id;
+ let local_index = self.bucket_to_global[bucket_id]
+ .iter()
+ .position(|&index| index == global_index)
+ .ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!("column {} missing from bucket {}",
global_index, bucket_id),
+ )
+ })?;
+ let state = self.bucket_states[bucket_id].as_ref().ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!("projected bucket {} was not loaded", bucket_id),
+ )
+ })?;
+ let encoded = match state {
+ BucketState::Monolithic { reader } =>
reader.encoded_column(local_index)?,
+ BucketState::Paged { column_readers } => column_readers
+ .get(local_index)
+ .and_then(Option::as_ref)
+ .ok_or_else(|| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!(
+ "projected column {} missing from paged bucket
{}",
+ global_index, bucket_id
+ ),
+ )
+ })?
+ .encoded_column(),
+ };
+ visitor(&column.name, &column.data_type, column.nullable,
encoded)?;
+ }
+ Ok(())
+ }
+
pub fn read_columns(&mut self) -> io::Result<RecordBatch> {
let num_cols = self.num_columns;
let mut arrays: Vec<Option<ArrayRef>> = vec![None; num_cols];
diff --git a/core/src/reader_tests.rs b/core/src/reader_tests.rs
index 87cdbb0..aa01139 100644
--- a/core/src/reader_tests.rs
+++ b/core/src/reader_tests.rs
@@ -16,6 +16,7 @@
// under the License.
use super::*;
+use crate::reader::{EncodedValueRef, Encoding};
use crate::writer::{MosaicWriter, OutputFile, WriterOptions};
use arrow_array::*;
use arrow_schema::{DataType, Field, Schema, TimeUnit};
@@ -103,6 +104,65 @@ fn values_to_batch(rows: &[Vec<Value>], columns:
&[(String, DataType, bool)]) ->
RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap()
}
+fn encoded_value_to_owned(value: EncodedValueRef<'_>) -> Value {
+ match value {
+ EncodedValueRef::Null => Value::Null,
+ EncodedValueRef::Boolean(value) => Value::Boolean(value),
+ EncodedValueRef::Int8(value) => Value::TinyInt(value),
+ EncodedValueRef::Int16(value) => Value::SmallInt(value),
+ EncodedValueRef::Int32(value) => Value::Integer(value),
+ EncodedValueRef::Int64(value) => Value::BigInt(value),
+ EncodedValueRef::Float32(value) => Value::Float(value),
+ EncodedValueRef::Float64(value) => Value::Double(value),
+ EncodedValueRef::Utf8(value) => Value::String(value.to_vec()),
+ EncodedValueRef::Binary(value) => Value::Bytes(value.to_vec()),
+ EncodedValueRef::DecimalCompact(value) => Value::DecimalCompact(value),
+ EncodedValueRef::DecimalLarge(value) =>
Value::DecimalLarge(value.to_vec()),
+ EncodedValueRef::Date32(value) => Value::Date(value),
+ EncodedValueRef::Time32(value) => Value::Time(value),
+ EncodedValueRef::TimestampMillis(value) =>
Value::TimestampMillis(value),
+ EncodedValueRef::TimestampMicros(value) =>
Value::TimestampMicros(value),
+ EncodedValueRef::TimestampNanos {
+ millis,
+ nanos_of_milli,
+ } => Value::TimestampNanos {
+ millis,
+ nanos_of_milli,
+ },
+ }
+}
+
+fn encoded_view_to_batch(row_group: &RowGroupReader) ->
io::Result<RecordBatch> {
+ let mut columns = Vec::new();
+ let mut values_by_column = Vec::new();
+ row_group.visit_encoded_columns(|name, data_type, nullable, column| {
+ columns.push((name.to_string(), data_type.clone(), nullable));
+ values_by_column.push(
+ column
+ .values()
+ .map(|value| value.map(encoded_value_to_owned))
+ .collect::<io::Result<Vec<_>>>()?,
+ );
+ Ok(())
+ })?;
+
+ let mut rows = vec![Vec::with_capacity(columns.len());
row_group.num_rows()];
+ for values in values_by_column {
+ for (row, value) in values.into_iter().enumerate() {
+ rows[row].push(value);
+ }
+ }
+ Ok(values_to_batch(&rows, &columns))
+}
+
+fn assert_encoded_view_matches_arrow(reader:
&MosaicReader<ByteArrayInputFile>) {
+ let mut arrow_row_group = reader.row_group_reader(0).unwrap();
+ let arrow = arrow_row_group.read_columns().unwrap();
+ let encoded_row_group = reader.row_group_reader(0).unwrap();
+ let encoded = encoded_view_to_batch(&encoded_row_group).unwrap();
+ assert_eq!(encoded, arrow);
+}
+
fn build_array_from_values(rows: &[Vec<Value>], col: usize, dt: &DataType) ->
Arc<dyn Array> {
match dt {
DataType::Boolean => {
@@ -2961,6 +3021,453 @@ fn write_and_read_paged(
(reader, data)
}
+fn encoded_view_fixture(page_size_threshold: usize) ->
MosaicReader<ByteArrayInputFile> {
+ let columns = vec![
+ ("plain".to_string(), DataType::Int32, true),
+ ("all_null".to_string(), DataType::Int64, true),
+ ("dict".to_string(), DataType::Utf8, true),
+ ("nullable_const".to_string(), DataType::Int16, true),
+ ];
+ let rows: Vec<Vec<Value>> = (0..200)
+ .map(|i| {
+ vec![
+ if i % 13 == 0 {
+ Value::Null
+ } else {
+ Value::Integer(i)
+ },
+ Value::Null,
+ if i % 10 == 0 {
+ Value::Null
+ } else {
+ Value::String(["red", "green", "blue"][i as usize %
3].as_bytes().to_vec())
+ },
+ if i % 4 == 0 {
+ Value::Null
+ } else {
+ Value::SmallInt(0)
+ },
+ ]
+ })
+ .collect();
+
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ &columns_to_arrow_schema(&columns),
+ WriterOptions {
+ num_buckets: 1,
+ page_size_threshold,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ write_values(&mut writer, &columns, &rows);
+ writer.close().unwrap();
+ let data = writer.output().buf.clone();
+ let len = data.len() as u64;
+ MosaicReader::new(ByteArrayInputFile::new(data), len).unwrap()
+}
+
+fn write_encoded_view_fixture(
+ columns: Vec<(String, DataType, bool)>,
+ rows: &[Vec<Value>],
+ page_size_threshold: usize,
+) -> MosaicReader<ByteArrayInputFile> {
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ &columns_to_arrow_schema(&columns),
+ WriterOptions {
+ num_buckets: 1,
+ page_size_threshold,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ write_values(&mut writer, &columns, rows);
+ writer.close().unwrap();
+ let data = writer.output().buf.clone();
+ MosaicReader::new(ByteArrayInputFile::new(data.clone()), data.len() as
u64).unwrap()
+}
+
+fn encoded_scalar_types() -> Vec<(&'static str, DataType)> {
+ vec![
+ ("bool", DataType::Boolean),
+ ("i8", DataType::Int8),
+ ("i16", DataType::Int16),
+ ("i32", DataType::Int32),
+ ("i64", DataType::Int64),
+ ("f32", DataType::Float32),
+ ("f64", DataType::Float64),
+ ("utf8", DataType::Utf8),
+ ("binary", DataType::Binary),
+ ("decimal_compact", DataType::Decimal128(18, 3)),
+ ("decimal_large", DataType::Decimal128(30, 4)),
+ ("date", DataType::Date32),
+ ("time", DataType::Time32(TimeUnit::Millisecond)),
+ (
+ "timestamp_millis",
+ DataType::Timestamp(TimeUnit::Millisecond, None),
+ ),
+ (
+ "timestamp_micros",
+ DataType::Timestamp(TimeUnit::Microsecond, None),
+ ),
+ (
+ "timestamp_nanos",
+ DataType::Timestamp(TimeUnit::Nanosecond, None),
+ ),
+ ("legacy_timestamp_nanos", legacy_timestamp_nanos_type()),
+ ]
+}
+
+fn encoded_scalar_value(data_type: &DataType, seed: usize) -> Value {
+ let seed_i64 = seed as i64;
+ match data_type {
+ DataType::Boolean => Value::Boolean(seed & 1 != 0),
+ DataType::Int8 => Value::TinyInt(seed as i8 - 48),
+ DataType::Int16 => Value::SmallInt(seed as i16 * 17 - 700),
+ DataType::Int32 => Value::Integer(seed as i32 * 10_007 - 50_000),
+ DataType::Int64 => Value::BigInt(seed_i64 * 1_000_003 - 9_000_000),
+ DataType::Float32 => Value::Float(seed as f32 * 1.25 - 20.5),
+ DataType::Float64 => Value::Double(seed as f64 * 0.125 - 7.75),
+ DataType::Utf8 =>
Value::String(format!("value-{seed:03}").into_bytes()),
+ DataType::Binary => Value::Bytes(vec![
+ seed as u8,
+ seed.wrapping_mul(17) as u8,
+ seed.wrapping_mul(31) as u8,
+ ]),
+ DataType::Decimal128(precision, _) if *precision <= 18 => {
+ Value::DecimalCompact(seed_i64 * 10_003 - 40_000)
+ }
+ DataType::Decimal128(_, _) => Value::DecimalLarge(
+ ((1i128 << 80) + seed as i128 * 10_007)
+ .to_be_bytes()
+ .to_vec(),
+ ),
+ DataType::Date32 => Value::Date(19_000 + seed as i32),
+ DataType::Time32(_) => Value::Time(1_000 + seed as i32 * 17),
+ DataType::Timestamp(TimeUnit::Millisecond, _) => {
+ Value::TimestampMillis(1_700_000_000_000 + seed_i64)
+ }
+ DataType::Timestamp(TimeUnit::Microsecond, _) => {
+ Value::TimestampMicros(1_700_000_000_000_000 + seed_i64)
+ }
+ DataType::Timestamp(TimeUnit::Nanosecond, _) | DataType::Struct(_)
+ if crate::types::is_timestamp_nanos(data_type) =>
+ {
+ Value::TimestampNanos {
+ millis: 1_700_000_000_000 + seed_i64,
+ nanos_of_milli: (seed * 7_919 % 1_000_000) as i32,
+ }
+ }
+ other => panic!("unsupported encoded scalar type: {other:?}"),
+ }
+}
+
+fn assert_fixture_encodings(
+ reader: &MosaicReader<ByteArrayInputFile>,
+ expected: &[(&str, Encoding)],
+) {
+ let infos = reader.page_infos(0).unwrap();
+ assert_eq!(infos.len(), expected.len());
+ for (name, encoding) in expected {
+ let column_index = reader
+ .schema()
+ .columns
+ .iter()
+ .position(|column| column.name == *name)
+ .unwrap();
+ let info = infos
+ .iter()
+ .find(|info| info.column_index == column_index)
+ .unwrap();
+ assert_eq!(info.encoding, *encoding, "column {name}");
+ }
+}
+
+fn assert_all_scalar_encoded_views(page_size_threshold: usize) {
+ let scalar_types = encoded_scalar_types();
+ let mut columns = Vec::new();
+ let mut expected = Vec::new();
+ for (name, data_type) in &scalar_types {
+ for (suffix, encoding) in [
+ ("all_null", Encoding::AllNull),
+ ("const", Encoding::Const),
+ ("dict", Encoding::Dict),
+ ] {
+ let column_name = format!("{name}_{suffix}");
+ columns.push((column_name.clone(), data_type.clone(), true));
+ expected.push((column_name, encoding));
+ }
+ if *data_type != DataType::Boolean {
+ let column_name = format!("{name}_plain");
+ columns.push((column_name.clone(), data_type.clone(), true));
+ expected.push((column_name, Encoding::Plain));
+ }
+ }
+
+ let rows = (0..96)
+ .map(|row| {
+ let mut values = Vec::with_capacity(columns.len());
+ for (_, data_type) in &scalar_types {
+ values.push(Value::Null);
+ values.push(if row % 5 == 0 {
+ Value::Null
+ } else {
+ encoded_scalar_value(data_type, 7)
+ });
+ values.push(if row % 11 == 0 {
+ Value::Null
+ } else {
+ encoded_scalar_value(data_type, row % 3)
+ });
+ if *data_type != DataType::Boolean {
+ values.push(if row % 17 == 0 {
+ Value::Null
+ } else {
+ encoded_scalar_value(data_type, row)
+ });
+ }
+ }
+ values
+ })
+ .collect::<Vec<_>>();
+
+ let reader = write_encoded_view_fixture(columns, &rows,
page_size_threshold);
+ let expected_refs = expected
+ .iter()
+ .map(|(name, encoding)| (name.as_str(), *encoding))
+ .collect::<Vec<_>>();
+ assert_fixture_encodings(&reader, &expected_refs);
+ assert_encoded_view_matches_arrow(&reader);
+
+ let bool_columns = vec![
+ ("bool_all_null".to_string(), DataType::Boolean, true),
+ ("bool_const".to_string(), DataType::Boolean, true),
+ ("bool_plain".to_string(), DataType::Boolean, true),
+ ("plain_anchor".to_string(), DataType::Int32, true),
+ ];
+ let bool_rows = vec![
+ vec![
+ Value::Null,
+ Value::Boolean(true),
+ Value::Boolean(true),
+ Value::Integer(1),
+ ],
+ vec![
+ Value::Null,
+ Value::Null,
+ Value::Boolean(false),
+ Value::Integer(2),
+ ],
+ ];
+ let bool_reader = write_encoded_view_fixture(bool_columns, &bool_rows,
page_size_threshold);
+ assert_fixture_encodings(
+ &bool_reader,
+ &[
+ ("bool_all_null", Encoding::AllNull),
+ ("bool_const", Encoding::Const),
+ ("bool_plain", Encoding::Plain),
+ ("plain_anchor", Encoding::Plain),
+ ],
+ );
+ assert_encoded_view_matches_arrow(&bool_reader);
+}
+
+fn assert_encoded_view(reader: &MosaicReader<ByteArrayInputFile>) {
+ assert_encoded_view_matches_arrow(reader);
+ let row_group = reader.row_group_reader(0).unwrap();
+ let mut names = Vec::new();
+ row_group
+ .visit_encoded_columns(|name, data_type, nullable, column| {
+ names.push(name.to_string());
+ assert_eq!(column.data_type(), data_type);
+ assert_eq!(column.num_rows(), 200);
+ assert!(nullable);
+
+ match name {
+ "plain" => {
+ assert_eq!(column.encoding(), Encoding::Plain);
+ assert!(column.has_nulls());
+ let null_bitmap = column.null_bitmap().unwrap();
+ assert_eq!(null_bitmap.len(), 25);
+ assert_eq!(column.constant().unwrap(), None);
+ for (row, value) in column.values().enumerate() {
+ let expected_null = row % 13 == 0;
+ let expected = if row % 13 == 0 {
+ EncodedValueRef::Null
+ } else {
+ EncodedValueRef::Int32(row as i32)
+ };
+ assert_eq!(value.unwrap(), expected);
+ assert_eq!(null_bitmap[row / 8] & (1 << (row % 8)) !=
0, expected_null);
+ assert_eq!(column.is_null(row), expected_null);
+ }
+ }
+ "all_null" => {
+ assert_eq!(column.encoding(), Encoding::AllNull);
+ assert!(column.has_nulls());
+ assert_eq!(column.null_bitmap(), None);
+ assert_eq!(column.constant().unwrap(), None);
+ for (row, value) in column.values().enumerate() {
+ assert_eq!(value.unwrap(), EncodedValueRef::Null);
+ assert!(column.is_null(row));
+ }
+ }
+ "dict" => {
+ assert_eq!(column.encoding(), Encoding::Dict);
+ assert!(column.has_nulls());
+ let null_bitmap = column.null_bitmap().unwrap();
+ assert_eq!(null_bitmap.len(), 25);
+ assert_eq!(column.constant().unwrap(), None);
+ for (row, value) in column.values().enumerate() {
+ let expected_null = row % 10 == 0;
+ let expected = if row % 10 == 0 {
+ EncodedValueRef::Null
+ } else {
+ EncodedValueRef::Utf8(["red", "green", "blue"][row
% 3].as_bytes())
+ };
+ assert_eq!(value.unwrap(), expected);
+ assert_eq!(null_bitmap[row / 8] & (1 << (row % 8)) !=
0, expected_null);
+ assert_eq!(column.is_null(row), expected_null);
+ }
+ }
+ "nullable_const" => {
+ assert_eq!(column.encoding(), Encoding::Const);
+ assert!(column.has_nulls());
+ let null_bitmap = column.null_bitmap().unwrap();
+ assert_eq!(null_bitmap.len(), 25);
+ assert_eq!(column.constant().unwrap(),
Some(EncodedValueRef::Int16(0)));
+ for (row, value) in column.values().enumerate() {
+ let expected_null = row % 4 == 0;
+ let expected = if row % 4 == 0 {
+ EncodedValueRef::Null
+ } else {
+ EncodedValueRef::Int16(0)
+ };
+ assert_eq!(value.unwrap(), expected);
+ assert_eq!(null_bitmap[row / 8] & (1 << (row % 8)) !=
0, expected_null);
+ assert_eq!(column.is_null(row), expected_null);
+ }
+ }
+ other => panic!("unexpected encoded column {other}"),
+ }
+ Ok(())
+ })
+ .unwrap();
+ assert_eq!(names, ["plain", "all_null", "dict", "nullable_const"]);
+
+ let projected = reader
+ .row_group_reader_by_names(0, &["dict", "nullable_const"])
+ .unwrap();
+ let mut projected_names = Vec::new();
+ projected
+ .visit_encoded_columns(|name, _, _, _| {
+ projected_names.push(name.to_string());
+ Ok(())
+ })
+ .unwrap();
+ assert_eq!(projected_names, ["dict", "nullable_const"]);
+
+ let duplicate = reader
+ .row_group_reader_by_names(0, &["dict", "dict"])
+ .unwrap();
+ let mut duplicate_names = Vec::new();
+ duplicate
+ .visit_encoded_columns(|name, _, _, _| {
+ duplicate_names.push(name.to_string());
+ Ok(())
+ })
+ .unwrap();
+ assert_eq!(duplicate_names, ["dict"]);
+
+ let mut duplicate_arrow = reader
+ .row_group_reader_by_names(0, &["dict", "dict"])
+ .unwrap();
+ let duplicate_batch = duplicate_arrow.read_columns().unwrap();
+ assert_eq!(duplicate_batch.num_columns(), 1);
+ assert_eq!(duplicate_batch.schema().field(0).name(), "dict");
+}
+
+#[test]
+fn test_visit_encoded_columns_monolithic_and_paged() {
+ assert_encoded_view(&encoded_view_fixture(usize::MAX));
+ assert_encoded_view(&encoded_view_fixture(1));
+}
+
+#[test]
+fn test_encoded_view_matches_arrow_for_all_scalar_types_and_encodings() {
+ assert_all_scalar_encoded_views(usize::MAX);
+ assert_all_scalar_encoded_views(1);
+}
+
+#[test]
+fn test_visit_encoded_columns_rejects_nested_before_callback() {
+ let item = Arc::new(Field::new("item", DataType::Int32, true));
+ let schema = Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("values", DataType::List(item), true),
+ ]);
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ &schema,
+ WriterOptions {
+ num_buckets: 1,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ let mut lists =
+
arrow_array::builder::ListBuilder::new(arrow_array::builder::Int32Builder::new());
+ for i in 0..4 {
+ lists.values().append_value(i);
+ lists.append(true);
+ }
+ let batch = RecordBatch::try_new(
+ Arc::new(schema),
+ vec![
+ Arc::new(Int32Array::from(vec![0, 1, 2, 3])),
+ Arc::new(lists.finish()),
+ ],
+ )
+ .unwrap();
+ writer.write_batch(&batch).unwrap();
+ writer.close().unwrap();
+ let data = writer.output().buf.clone();
+ let reader =
+ MosaicReader::new(ByteArrayInputFile::new(data.clone()), data.len() as
u64).unwrap();
+
+ let row_group = reader.row_group_reader(0).unwrap();
+ let mut callbacks = 0;
+ let error = row_group
+ .visit_encoded_columns(|_, _, _, _| {
+ callbacks += 1;
+ Ok(())
+ })
+ .unwrap_err();
+ assert_eq!(error.kind(), io::ErrorKind::Unsupported);
+ assert!(error.to_string().contains("values"));
+ assert_eq!(callbacks, 0);
+
+ let id = reader
+ .schema()
+ .columns
+ .iter()
+ .position(|column| column.name == "id")
+ .unwrap();
+ let projected = reader.row_group_reader_projected(0, &[id]).unwrap();
+ projected
+ .visit_encoded_columns(|name, _, _, column| {
+ assert_eq!(name, "id");
+ assert_eq!(column.encoding(), Encoding::Plain);
+ Ok(())
+ })
+ .unwrap();
+}
+
#[test]
fn test_paged_roundtrip_basic() {
let columns = vec![
@@ -4636,3 +5143,23 @@ fn test_slot_sizes_paged_array_column() {
assert_eq!(projected[0].column_index, vals);
assert_eq!(projected[0].slot_size, sizes[vals]);
}
+
+#[test]
+fn test_unknown_paged_encoding_returns_invalid_data_without_panicking() {
+ let result = std::panic::catch_unwind(|| -> io::Result<()> {
+ let page =
MosaicReader::<ByteArrayInputFile>::parse_simple_column_slot(
+ vec![0xff, 0x01],
+ &DataType::Int32,
+ 1,
+ )?;
+ let column = page.encoded_column();
+ column.values().next().transpose()?;
+ Ok(())
+ });
+
+ let err = result
+ .expect("unknown paged encoding must not panic")
+ .expect_err("unknown paged encoding must be rejected");
+ assert_eq!(err.kind(), io::ErrorKind::InvalidData);
+ assert!(err.to_string().contains("unsupported encoding 255"));
+}
