use std::{collections::BTreeMap, sync::Arc};
use lance_core::{Error, Result, datatypes::Schema};
use lance_encoding::{decoder::ColumnInfo, format::pb21};
use crate::{
format::pbfile,
reader::{
BufferDescriptor, CachedFileMetadata, FileMetadataIndex, RawFileMetadata, ReaderProjection,
structural,
},
version::ConcreteFileVersion,
};
fn required<'a, T>(value: Option<&'a T>, label: &str) -> Result<&'a T> {
value.ok_or_else(|| {
Error::invalid_input_source(
format!("Lance v2.1 {label} is missing its nested encoding").into(),
)
})
}
fn validate_compressive_encoding(encoding: &pb21::CompressiveEncoding) -> Result<()> {
use pb21::compressive_encoding::Compression;
match encoding.compression.as_ref() {
Some(Compression::Flat(_))
| Some(Compression::InlineBitpacking(_))
| Some(Compression::Constant(_)) => Ok(()),
Some(Compression::Variable(variable)) => validate_compressive_encoding(required(
variable.offsets.as_deref(),
"variable offsets",
)?),
Some(Compression::OutOfLineBitpacking(bitpacking)) => {
validate_compressive_encoding(required(
bitpacking.values.as_deref(),
"out-of-line bitpacking values",
)?)
}
Some(Compression::Fsst(fsst)) => {
validate_compressive_encoding(required(fsst.values.as_deref(), "FSST values")?)
}
Some(Compression::Dictionary(dictionary)) => {
validate_compressive_encoding(required(
dictionary.indices.as_deref(),
"dictionary indices",
)?)?;
validate_compressive_encoding(required(
dictionary.items.as_deref(),
"dictionary items",
)?)
}
Some(Compression::Rle(rle)) => {
let values = required(rle.values.as_deref(), "RLE values")?;
let run_lengths = required(rle.run_lengths.as_deref(), "RLE run lengths")?;
let fixed_values = matches!(
values.compression.as_ref(),
Some(Compression::Flat(flat))
if matches!(flat.bits_per_value, 8 | 16 | 32 | 64)
&& flat.data.is_none()
);
let fixed_u8_lengths = matches!(
run_lengths.compression.as_ref(),
Some(Compression::Flat(flat))
if flat.bits_per_value == 8 && flat.data.is_none()
);
if !fixed_values || !fixed_u8_lengths {
return Err(Error::invalid_input_source(
"Lance v2.1 RLE requires flat values and flat u8 run lengths".into(),
));
}
Ok(())
}
Some(Compression::ByteStreamSplit(split)) => validate_compressive_encoding(required(
split.values.as_deref(),
"byte-stream-split values",
)?),
Some(Compression::General(general)) => validate_compressive_encoding(required(
general.values.as_deref(),
"general-compression values",
)?),
Some(Compression::FixedSizeList(list)) => validate_compressive_encoding(required(
list.values.as_deref(),
"fixed-size-list values",
)?),
Some(Compression::PackedStruct(packed)) => validate_compressive_encoding(required(
packed.values.as_deref(),
"packed-struct values",
)?),
Some(Compression::VariablePackedStruct(_)) => Err(Error::invalid_input_source(
"Variable packed struct compression is not part of the Lance v2.1 grammar".into(),
)),
None => Err(Error::invalid_input_source(
"Lance v2.1 compressive encoding is missing its compression variant".into(),
)),
}
}
fn validate_page_layout(layout: &pb21::PageLayout) -> Result<()> {
use pb21::page_layout::Layout;
match layout.layout.as_ref() {
Some(Layout::MiniBlockLayout(miniblock)) => {
if miniblock.has_large_chunk {
return Err(Error::invalid_input_source(
"Large miniblock chunks are not part of the Lance v2.1 grammar".into(),
));
}
if let Some(rep) = miniblock.rep_compression.as_ref() {
validate_compressive_encoding(rep)?;
}
if let Some(def) = miniblock.def_compression.as_ref() {
validate_compressive_encoding(def)?;
}
validate_compressive_encoding(required(
miniblock.value_compression.as_ref(),
"miniblock values",
)?)?;
if let Some(dictionary) = miniblock.dictionary.as_ref() {
validate_compressive_encoding(dictionary)?;
}
Ok(())
}
Some(Layout::FullZipLayout(fullzip)) => validate_compressive_encoding(required(
fullzip.value_compression.as_ref(),
"full-zip values",
)?),
Some(Layout::ConstantLayout(constant)) => {
if constant.inline_value.is_some() {
Err(Error::invalid_input_source(
"Lance v2.1 only accepts the all-null form of constant page layout".into(),
))
} else if constant.rep_compression.is_some() || constant.def_compression.is_some() {
Err(Error::invalid_input_source(
"Compressed constant-page levels are not part of the Lance v2.1 grammar".into(),
))
} else {
Ok(())
}
}
Some(Layout::BlobLayout(blob)) => {
let inner = blob.inner_layout.as_deref().ok_or_else(|| {
Error::invalid_input_source(
"Lance v2.1 blob page layout is missing its inner layout".into(),
)
})?;
validate_page_layout(inner)
}
Some(Layout::SparseLayout(_)) => Err(Error::invalid_input_source(
"Sparse page layout is not part of the Lance v2.1 grammar".into(),
)),
None => Err(Error::invalid_input_source(
"Lance v2.1 page is missing its page layout".into(),
)),
}
}
pub fn decode_column(
column_index: u32,
metadata: &pbfile::ColumnMetadata,
) -> Result<Arc<ColumnInfo>> {
let page_infos = metadata
.pages
.iter()
.enumerate()
.map(|(page_index, page)| {
let page_layout = structural::decode_page_layout(column_index, page_index, page)?;
validate_page_layout(&page_layout)?;
structural::build_page_info(column_index, page_index, page, page_layout)
})
.collect::<Result<Vec<_>>>()?;
structural::build_column_info(column_index, metadata, page_infos)
}
pub fn decode_column_metadata(
column_metadatas: &[pbfile::ColumnMetadata],
) -> Result<Vec<Arc<ColumnInfo>>> {
column_metadatas
.iter()
.enumerate()
.map(|(column_index, metadata)| {
let column_index = u32::try_from(column_index).map_err(|_| {
Error::invalid_input_source("File has more than u32::MAX columns".into())
})?;
decode_column(column_index, metadata)
})
.collect()
}
pub fn projection_from_field_ids(
schema: &Schema,
field_id_to_column_index: &BTreeMap<u32, u32>,
) -> ReaderProjection {
structural::projection_from_field_ids(schema, field_id_to_column_index)
}
pub fn projection_from_whole_schema(schema: &Schema) -> ReaderProjection {
structural::projection_from_field_ids(schema, &super::field_id_to_column_index(schema))
}
pub fn projection_from_column_names(
schema: &Schema,
column_names: &[&str],
) -> Result<ReaderProjection> {
structural::projection_from_column_names(
schema,
column_names,
&super::field_id_to_column_index(schema),
)
}
pub fn finish_metadata(raw: RawFileMetadata) -> Result<CachedFileMetadata> {
if (raw.footer.major_version, raw.footer.minor_version) != (2, 1) {
return Err(Error::version_conflict(
"Attempt to use the Lance v2.1 reader for a different file version".to_string(),
raw.footer.major_version,
raw.footer.minor_version,
));
}
validate_global_buffers(&raw.file_buffers)?;
let column_infos = decode_column_metadata(&raw.column_metadatas)?;
Ok(CachedFileMetadata {
file_schema: raw.file_schema,
column_metadatas: raw.column_metadatas,
column_infos,
num_rows: raw.num_rows,
file_buffers: raw.file_buffers,
num_data_bytes: raw.num_data_bytes,
num_column_metadata_bytes: raw.num_column_metadata_bytes,
num_global_buffer_bytes: raw.num_global_buffer_bytes,
num_footer_bytes: raw.num_footer_bytes,
major_version: raw.footer.major_version,
minor_version: raw.footer.minor_version,
version: ConcreteFileVersion::V2_1,
file_size_bytes: raw.file_size_bytes,
retained_global_buffers: raw.retained_global_buffers,
})
}
pub fn validate_global_buffers(buffers: &[BufferDescriptor]) -> Result<()> {
structural::validate_global_buffers(buffers)
}
pub fn finish_metadata_index(index: FileMetadataIndex) -> Result<FileMetadataIndex> {
if index.version != ConcreteFileVersion::V2_1 {
let (major, minor) = index.version.to_standard_footer_numbers();
return Err(Error::version_conflict(
"Attempt to use the Lance v2.1 reader for a different metadata index".to_string(),
major,
minor,
));
}
validate_global_buffers(&index.file_buffers)?;
Ok(index)
}
#[cfg(test)]
pub fn test_projection_length(
schema: &Schema,
column_indices: &[u32],
column_lengths: &[u64],
) -> Result<u64> {
structural::test_projection_length(schema, column_indices, column_lengths)
}
#[cfg(test)]
mod grammar_tests {
use super::*;
use pb21::{
CompressiveEncoding, FullZipLayout, General, PageLayout, VariablePackedStruct,
compressive_encoding::Compression, page_layout::Layout,
};
#[test]
fn rejects_nested_variable_packed_struct() {
let variable_packed = CompressiveEncoding {
compression: Some(Compression::VariablePackedStruct(VariablePackedStruct {
fields: Vec::new(),
})),
};
let wrapped = CompressiveEncoding {
compression: Some(Compression::General(Box::new(General {
compression: None,
values: Some(Box::new(variable_packed)),
}))),
};
let layout = PageLayout {
layout: Some(Layout::FullZipLayout(FullZipLayout {
value_compression: Some(wrapped),
..Default::default()
})),
};
let error = validate_page_layout(&layout).unwrap_err();
assert!(
error
.to_string()
.contains("Variable packed struct compression is not part")
);
}
}