diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 46cb7f499d3..fb1de8addcf 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -4307,7 +4307,9 @@ def _assert_mixed_version_reads(dataset, expected): ("2.0", "2.1", "2.2", "2.3"), ("2.1", "2.2", "2.3", "2.0"), ("2.2", "2.3", "2.0", "2.1"), - ("2.3", "2.0", "2.1", "2.2"), + # A table created on 2.3 follows the semantic type contract and only + # takes data files of version 2.1 or later. + ("2.3", "2.1", "2.2", "2.1"), ], ) @pytest.mark.parametrize("mode", ["reencode", "try_binary_copy"]) diff --git a/rust/lance-arrow/src/scalar.rs b/rust/lance-arrow/src/scalar.rs index 6003eb13497..c9271e5926f 100644 --- a/rust/lance-arrow/src/scalar.rs +++ b/rust/lance-arrow/src/scalar.rs @@ -94,6 +94,15 @@ pub fn encode_scalar_value_buffer(scalar: &ArrayRef) -> Result> { Ok(out) } +/// The lengths of the Arrow buffers a scalar value buffer holds, in order. +pub fn scalar_value_buffer_lengths(value_buffer: &[u8]) -> Result> { + let mut offset = 0; + let num_buffers = read_u32(value_buffer, &mut offset)? as usize; + (0..num_buffers) + .map(|_| read_u32(value_buffer, &mut offset).map(|len| len as usize)) + .collect() +} + pub fn decode_scalar_from_value_buffer( data_type: &DataType, value_buffer: &[u8], diff --git a/rust/lance-core/src/datatypes/field.rs b/rust/lance-core/src/datatypes/field.rs index efbc3401d33..c1b54df124f 100644 --- a/rust/lance-core/src/datatypes/field.rs +++ b/rust/lance-core/src/datatypes/field.rs @@ -25,7 +25,8 @@ use lance_arrow::{ }; use super::{ - Dictionary, LogicalType, OUTPUT_ENCODING_META_KEY, OutputEncoding, Projection, TypeComparison, + Dictionary, LogicalType, OUTPUT_ENCODING_META_KEY, OutputEncoding, Projection, SemanticType, + SemanticTypeClass, TypeComparison, schema::{compare_fields, explain_fields_difference}, }; use crate::{ @@ -160,6 +161,16 @@ pub struct Field { /// None means the field is not part of the clustering key. /// Some(n) means this field is the nth column in the clustering key. pub unenforced_clustering_key_position: Option, + + /// The Arrow layout of this field under the semantic type contract. + /// + /// Fields of a table that follows the contract always carry it when their + /// semantic type has output encodings: it is the table's output encoding, + /// or the layout of data being written to the table. Legacy table and data + /// file fields never carry it, because their `logical_type` names exactly + /// one Arrow type. A table records it as the [`OUTPUT_ENCODING_META_KEY`] + /// field metadata entry when it is not the type's default. + pub output_encoding: Option, } impl Field { @@ -171,6 +182,9 @@ impl Field { /// Returns arrow data type. pub fn data_type(&self) -> DataType { + if let Some(data_type) = self.output_encoding_data_type() { + return data_type; + } match &self.logical_type { lt if lt.is_list() => DataType::List(Arc::new(ArrowField::from(&self.children[0]))), lt if lt.is_large_list() => { @@ -196,6 +210,44 @@ impl Field { } } + /// The Arrow layout [`Self::output_encoding`] selects, when it selects one. + /// + /// Output encodings of value-transforming types select a conversion rather + /// than a layout, so those fields keep the type `logical_type` names. + fn output_encoding_data_type(&self) -> Option { + let encoding = self.output_encoding.as_ref()?; + let semantic = self.logical_type.semantic().ok()?; + let item = self.children.first().map(ArrowField::from); + semantic + .semantic_type + .layout_data_type(encoding, item.as_ref()) + .ok() + } + + /// The Arrow type of `other` that projections and intersections compare + /// with this field's. + /// + /// A field with an output encoding follows the semantic type contract, + /// where another layout of the same semantic type holds the same values, + /// so it compares as this field's own layout. Otherwise it is `other`'s. + fn layout_to_compare(&self, other: &Self) -> DataType { + if self.output_encoding.is_some() && self.type_matches(other, TypeComparison::Semantic) { + self.data_type() + } else { + other.data_type() + } + } + + /// This field's semantic type and the output encoding of its Arrow layout. + fn semantic_layout(&self) -> Option<(SemanticType, Option)> { + let semantic = self.logical_type.semantic().ok()?; + let encoding = self + .output_encoding + .clone() + .or_else(|| semantic.output_encoding()); + Some((semantic.semantic_type, encoding)) + } + pub fn has_dictionary_types(&self) -> bool { matches!(self.data_type(), DataType::Dictionary(_, _)) || self.children.iter().any(Self::has_dictionary_types) @@ -445,18 +497,29 @@ impl Field { /// Whether this field's own type is compatible with `expected`'s under /// `comparison`. Children are not compared. pub fn type_matches(&self, expected: &Self, comparison: TypeComparison) -> bool { - comparison.logical_types_match(&self.logical_type, &expected.logical_type) + if self.output_encoding.is_none() && expected.output_encoding.is_none() { + return comparison.logical_types_match(&self.logical_type, &expected.logical_type); + } + // With an output encoding the string no longer names the layout, so an + // exact comparison compares the layouts it resolves to. + match (self.semantic_layout(), expected.semantic_layout()) { + (Some((actual_type, actual_layout)), Some((expected_type, expected_layout))) => { + actual_type == expected_type + && (comparison == TypeComparison::Semantic || actual_layout == expected_layout) + } + _ => false, + } } - /// This field as a table that follows the semantic type contract records - /// it. + /// This field as a table that follows the semantic type contract holds it. /// - /// A legacy alias becomes its canonical semantic type, and the layout the - /// alias named becomes the field's [`OUTPUT_ENCODING_META_KEY`] entry when - /// it is not the type's default, so reads return the same Arrow type. An - /// existing entry takes precedence over the implied layout and is kept - /// unchanged, including values this build does not recognize. Children are - /// converted the same way; everything else is preserved. + /// A legacy alias becomes its canonical semantic type. The output encoding + /// comes from, in order, a valid [`OUTPUT_ENCODING_META_KEY`] entry, which + /// moves out of the metadata, the output encoding already set, the layout + /// the alias named, and the type's default, so reads return the same Arrow + /// type. An entry this build does not recognize, or that is not valid for + /// the type, stays in the metadata unchanged and has no effect. Children + /// are converted the same way, and everything else is preserved. pub fn to_canonical_type(&self) -> Result { let semantic = self.logical_type.semantic().map_err(|err| { Error::schema(format!( @@ -465,13 +528,14 @@ impl Field { )) })?; let mut field = self.clone(); - field.logical_type = semantic.semantic_type.logical_type(); - if let Some(implied) = semantic.implied_encoding { - field - .metadata - .entry(OUTPUT_ENCODING_META_KEY.to_string()) - .or_insert_with(|| implied.to_string()); + let recorded = self.recorded_output_encoding(); + if recorded.is_some() { + field.metadata.remove(OUTPUT_ENCODING_META_KEY); } + field.output_encoding = recorded + .or_else(|| self.output_encoding.clone()) + .or_else(|| semantic.output_encoding()); + field.logical_type = semantic.semantic_type.logical_type(); field.children = self .children .iter() @@ -480,6 +544,123 @@ impl Field { Ok(field) } + /// This field as a data file schema records it: the Arrow-mapped + /// `logical_type` of its exact layout, without an output encoding. + /// + /// A view layout is recorded as the offset layout with 64-bit offsets, + /// which holds every value; writers convert view arrays to it. Fields + /// whose `logical_type` already names their layout are unchanged. + pub fn to_data_file_field(&self) -> Result { + let mut field = self.clone(); + if self.is_representation_only() { + let data_type = match self.data_type() { + DataType::Utf8View => DataType::LargeUtf8, + DataType::BinaryView => DataType::LargeBinary, + data_type => data_type, + }; + field.logical_type = LogicalType::try_from(&data_type)?; + } + field.output_encoding = None; + field.children = self + .children + .iter() + .map(Self::to_data_file_field) + .collect::>()?; + Ok(field) + } + + /// Record the view layouts of `input`, the Arrow field of data written to + /// this field, as output encodings. + /// + /// Converting from Arrow reads a view as the offset layout of the same + /// width, which is how legacy tables store it. Under the semantic type + /// contract a view is its own output encoding, and writers store view data + /// in a layout that holds every value. + pub fn record_view_layouts(&mut self, input: &ArrowField) { + if matches!(input.data_type(), DataType::Utf8View | DataType::BinaryView) { + self.output_encoding = OutputEncoding::of_data_type(input.data_type()); + } + match input.data_type() { + DataType::Struct(input_children) => { + for input_child in input_children { + if let Some(child) = self.child_mut(input_child.name()) { + child.record_view_layouts(input_child); + } + } + } + DataType::List(item) + | DataType::LargeList(item) + | DataType::FixedSizeList(item, _) + | DataType::Map(item, _) => { + if let Some(child) = self.children.first_mut() { + child.record_view_layouts(item); + } + } + _ => {} + } + } + + fn is_representation_only(&self) -> bool { + self.logical_type.semantic().is_ok_and(|semantic| { + semantic.semantic_type.class() == SemanticTypeClass::RepresentationOnly + }) + } + + /// This table field with the Arrow layouts of `input`, a field of data + /// being written to it. + /// + /// Representation-only fields take `input`'s layout, so each data file + /// records the layout it actually encodes. Everything else, including field + /// IDs and metadata, comes from this field. Struct children are matched by + /// name and the child of a list or map by position; children missing from + /// `input` are left out. The types must already be compatible under + /// [`TypeComparison::Semantic`]. + pub fn with_input_layout(&self, input: &Self) -> Result { + let mut field = self.clone(); + if self.is_blob() { + return Ok(field); + } + if self.is_representation_only() { + field.logical_type = input.logical_type.clone(); + field.output_encoding = input.output_encoding.clone(); + } + field.children = if self.logical_type.is_struct() { + input + .children + .iter() + .map(|input_child| { + self.child(&input_child.name) + .ok_or_else(|| { + Error::schema(format!( + "field '{}' has no child '{}'", + self.name, input_child.name + )) + })? + .with_input_layout(input_child) + }) + .collect::>()? + } else { + self.children + .iter() + .zip(&input.children) + .map(|(child, input_child)| child.with_input_layout(input_child)) + .collect::>()? + }; + Ok(field) + } + + /// The value a table records in this field's [`OUTPUT_ENCODING_META_KEY`] + /// entry: its output encoding, unless that is the type's default. + pub fn output_encoding_entry(&self) -> Option { + let encoding = self.output_encoding.as_ref()?; + let default = self + .logical_type + .semantic() + .ok() + .and_then(|semantic| semantic.semantic_type.default_output_encoding()); + (default.as_ref() != Some(encoding)).then(|| encoding.to_string()) + } + /// The output encoding named by this field's [`OUTPUT_ENCODING_META_KEY`] /// entry, if it is present and valid for the field's semantic type. /// @@ -766,6 +947,7 @@ impl Field { dictionary: self.dictionary.clone(), unenforced_primary_key_position: self.unenforced_primary_key_position, unenforced_clustering_key_position: self.unenforced_clustering_key_position, + output_encoding: self.output_encoding.clone(), }; if path_components.is_empty() { // Project stops here, copy all the remaining children. @@ -856,7 +1038,7 @@ impl Field { return Ok(self.clone()); } - match (self.data_type(), other.data_type()) { + match (self.data_type(), self.layout_to_compare(other)) { (DataType::Boolean, DataType::Boolean) => Ok(self.clone()), (dt, other_dt) if (dt.is_primitive() && other_dt.is_primitive()) @@ -1000,7 +1182,7 @@ impl Field { } let self_type = self.data_type(); - let other_type = other.data_type(); + let other_type = self.layout_to_compare(other); if matches!( (&self_type, &other_type), @@ -1039,6 +1221,7 @@ impl Field { dictionary: self.dictionary.clone(), unenforced_primary_key_position: self.unenforced_primary_key_position, unenforced_clustering_key_position: self.unenforced_clustering_key_position, + output_encoding: self.output_encoding.clone(), }; return Ok(f); } @@ -1100,6 +1283,7 @@ impl Field { dictionary: self.dictionary.clone(), unenforced_primary_key_position: self.unenforced_primary_key_position, unenforced_clustering_key_position: self.unenforced_clustering_key_position, + output_encoding: self.output_encoding.clone(), }) } } @@ -1371,6 +1555,7 @@ impl TryFrom<&ArrowField> for Field { dictionary: None, unenforced_primary_key_position, unenforced_clustering_key_position, + output_encoding: None, }) } } @@ -1688,9 +1873,10 @@ mod tests { ), true, ); - let mut schema = - crate::datatypes::Schema::try_from(&arrow_schema::Schema::new(vec![arrow_field])) - .unwrap(); + let mut schema = crate::datatypes::Schema::try_from(&arrow_schema::Schema::new(vec![ + arrow_field.clone(), + ])) + .unwrap(); schema.set_field_id(None); let canonical = schema.to_canonical_types().unwrap(); @@ -1700,7 +1886,7 @@ mod tests { ( field.id, field.logical_type.to_string(), - field.metadata.get(OUTPUT_ENCODING_META_KEY).cloned(), + field.output_encoding.as_ref().map(ToString::to_string), field.nullable, ) }) @@ -1718,19 +1904,25 @@ mod tests { encoding("dictionary:int8:utf8"), true ), - (4, "decimal:10:2".to_string(), None, true), + (4, "decimal:10:2".to_string(), encoding("decimal128"), true), (5, "decimal:10:2".to_string(), encoding("decimal256"), true), (6, "int64".to_string(), None, false), ] ); let name = canonical.field("s.name").unwrap(); assert_eq!(name.metadata.get("user").map(String::as_str), Some("kept")); + // Only encodings other than the default are recorded. + assert_eq!(name.output_encoding_entry().as_deref(), Some("large_utf8")); assert_eq!( - name.recorded_output_encoding(), - Some(OutputEncoding::LargeUtf8) + canonical.field("s.price").unwrap().output_encoding_entry(), + None ); + // Every field reads back as the Arrow type it was created from. + assert_eq!(ArrowField::from(&canonical.fields[0]), arrow_field); // The transform is idempotent, so a canonical schema is its own form. assert_eq!(canonical.to_canonical_types().unwrap(), canonical); + // The data file form names the exact layouts again. + assert_eq!(canonical.to_data_file_schema().unwrap(), schema); } #[test] @@ -1742,9 +1934,22 @@ mod tests { ); let canonical = field.to_canonical_type().unwrap(); assert_eq!(canonical.logical_type.to_string(), "string"); + assert_eq!(canonical.output_encoding, Some(OutputEncoding::Utf8View)); + assert!(!canonical.metadata.contains_key(OUTPUT_ENCODING_META_KEY)); + assert_eq!(canonical.data_type(), DataType::Utf8View); + + // An entry a reader cannot apply stays in the metadata, unused. + field + .metadata + .insert(OUTPUT_ENCODING_META_KEY.to_string(), "utf16".to_string()); + let canonical = field.to_canonical_type().unwrap(); + assert_eq!(canonical.output_encoding, Some(OutputEncoding::LargeUtf8)); assert_eq!( - canonical.recorded_output_encoding(), - Some(OutputEncoding::Utf8View) + canonical + .metadata + .get(OUTPUT_ENCODING_META_KEY) + .map(String::as_str), + Some("utf16") ); let dictionary_of_integers = Field::new_arrow( @@ -1779,6 +1984,92 @@ mod tests { assert_eq!(field.recorded_output_encoding(), expected); } + /// Views are stored as the offset layout that holds every value, and + /// canonical-only names become the layout they read as. + #[rstest::rstest] + #[case::utf8_view(DataType::Utf8View, "large_string")] + #[case::binary_view(DataType::BinaryView, "large_binary")] + #[case::large_utf8(DataType::LargeUtf8, "large_string")] + #[case::decimal(DataType::Decimal128(10, 2), "decimal:128:10:2")] + #[case::wide_decimal(DataType::Decimal256(40, 2), "decimal:256:40:2")] + fn test_to_data_file_field(#[case] layout: DataType, #[case] recorded: &str) { + let mut field = Field::new_arrow("a", layout.clone(), true) + .unwrap() + .to_canonical_type() + .unwrap(); + field.record_view_layouts(&ArrowField::new("a", layout, true)); + let data_file_field = field.to_data_file_field().unwrap(); + assert_eq!(data_file_field.logical_type.to_string(), recorded); + assert_eq!(data_file_field.output_encoding, None); + } + + /// A write takes field IDs and metadata from the table and layouts from + /// the input. + #[test] + fn test_input_layouts() { + let struct_of = |data_type: DataType| { + DataType::Struct(vec![ArrowField::new("b", data_type, true)].into()) + }; + let mut table = Field::new_arrow("a", struct_of(DataType::Utf8View), true).unwrap(); + table.record_view_layouts(&ArrowField::new("a", struct_of(DataType::Utf8View), true)); + let mut id = 0; + table.set_id(-1, &mut id); + table.children[0] + .metadata + .insert("lance-encoding:compression".to_string(), "zstd".to_string()); + let table = table.to_canonical_type().unwrap(); + assert_eq!( + table.children[0].output_encoding, + Some(OutputEncoding::Utf8View) + ); + + let input = Field::new_arrow("a", struct_of(DataType::LargeUtf8), true).unwrap(); + assert!(!input.compare_with_options(&table, &SchemaCompareOptions::default())); + let write = table.with_input_layout(&input).unwrap(); + assert_eq!(write.children[0].data_type(), DataType::LargeUtf8); + assert_eq!(write.children[0].id, 1); + assert_eq!( + write.children[0].metadata.get("lance-encoding:compression"), + Some(&"zstd".to_string()) + ); + + // Exact comparisons of fields with output encodings compare layouts. + assert!(!write.children[0].type_matches(&table.children[0], TypeComparison::Exact)); + assert!(write.children[0].type_matches(&table.children[0], TypeComparison::Semantic)); + } + + #[test] + fn test_check_output_encoding_entries() { + let schema_with = |value: &str| { + let mut field = Field::new_arrow("a", DataType::Utf8, true).unwrap(); + field.id = 0; + field + .metadata + .insert(OUTPUT_ENCODING_META_KEY.to_string(), value.to_string()); + crate::datatypes::Schema { + fields: vec![field], + metadata: HashMap::new(), + } + }; + schema_with("large_utf8") + .check_output_encoding_entries(None) + .unwrap(); + let err = schema_with("decimal128") + .check_output_encoding_entries(None) + .unwrap_err(); + assert!(err.to_string().contains("decimal128"), "{err}"); + // An entry carried unchanged was not set by this update. + let unknown = schema_with("utf16"); + unknown + .check_output_encoding_entries(Some(&unknown)) + .unwrap(); + assert!( + schema_with("utf16") + .check_output_encoding_entries(Some(&schema_with("utf32"))) + .is_err() + ); + } + #[test] fn test_compare_type_comparison() { let table = Field::new_arrow( diff --git a/rust/lance-core/src/datatypes/schema.rs b/rust/lance-core/src/datatypes/schema.rs index 5fec2383413..865a0bb12f3 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -14,6 +14,7 @@ use arrow_array::RecordBatch; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use lance_arrow::*; +use super::OUTPUT_ENCODING_META_KEY; use super::field::{Field, OnTypeMismatch, SchemaCompareOptions}; use crate::{ Error, ROW_ADDR, ROW_ADDR_FIELD, ROW_CREATED_AT_VERSION, ROW_CREATED_AT_VERSION_FIELD, ROW_ID, @@ -215,6 +216,102 @@ impl Schema { }) } + /// Reject [`OUTPUT_ENCODING_META_KEY`] entries that this schema sets to a + /// value that is unknown or invalid for the field's semantic type. + /// + /// An entry that `previous` already holds with the same value, for the + /// field with the same ID, is not being set and passes: readers ignore + /// such entries, and a newer build may have written them. + pub fn check_output_encoding_entries(&self, previous: Option<&Self>) -> Result<()> { + for field in self.fields_pre_order() { + let Some(value) = field.metadata.get(OUTPUT_ENCODING_META_KEY) else { + continue; + }; + if field.recorded_output_encoding().is_some() { + continue; + } + let unchanged = previous + .and_then(|previous| previous.field_by_id(field.id)) + .and_then(|previous| previous.metadata.get(OUTPUT_ENCODING_META_KEY)) + == Some(value); + if !unchanged { + return Err(Error::invalid_input(format!( + "cannot set {OUTPUT_ENCODING_META_KEY} of field '{}' to '{value}': it is not an output encoding of type '{}'", + field.name, field.logical_type + ))); + } + } + Ok(()) + } + + /// This schema as a data file schema records it. See + /// [`Field::to_data_file_field`]. + pub fn to_data_file_schema(&self) -> Result { + Ok(Self { + fields: self + .fields + .iter() + .map(Field::to_data_file_field) + .collect::>()?, + metadata: self.metadata.clone(), + }) + } + + /// Record the view layouts of `input`, the Arrow schema of data written + /// with this schema. See [`Field::record_view_layouts`]. + pub fn record_view_layouts(&mut self, input: &ArrowSchema) { + for input_field in input.fields() { + if let Some(field) = self + .fields + .iter_mut() + .find(|f| &f.name == input_field.name()) + { + field.record_view_layouts(input_field); + } + } + } + + /// This write schema with the Arrow layouts of `input`, the schema of the + /// data being written. Fields `input` does not hold keep their layout. + /// See [`Field::with_input_layout`]. + pub fn with_input_layouts(&self, input: &Self) -> Result { + let fields = self + .fields + .iter() + .map(|field| match input.field(&field.name) { + Some(input_field) => field.with_input_layout(input_field), + None => Ok(field.clone()), + }) + .collect::>()?; + Ok(Self { + fields, + metadata: self.metadata.clone(), + }) + } + + /// The fields of this table schema that `input` writes, with `input`'s + /// Arrow layouts. See [`Field::with_input_layout`]. + pub fn project_for_write(&self, input: &Self) -> Result { + let fields = input + .fields + .iter() + .map(|input_field| { + self.field(&input_field.name) + .ok_or_else(|| { + Error::schema(format!( + "field '{}' does not exist in the table schema", + input_field.name + )) + })? + .with_input_layout(input_field) + }) + .collect::>()?; + Ok(Self { + fields, + metadata: self.metadata.clone(), + }) + } + pub fn check_compatible(&self, expected: &Self, options: &SchemaCompareOptions) -> Result<()> { if !self.compare_with_options(expected, options) { let difference = self.explain_difference(expected, options); diff --git a/rust/lance-core/src/datatypes/semantic.rs b/rust/lance-core/src/datatypes/semantic.rs index 9d83ee3f3a3..eaeea59fe22 100644 --- a/rust/lance-core/src/datatypes/semantic.rs +++ b/rust/lance-core/src/datatypes/semantic.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use arrow_schema::{DataType, Field as ArrowField}; use super::LogicalType; +use crate::deepsize::DeepSizeOf; use crate::{Error, Result}; /// Field metadata entry naming the Arrow layout that reads return for a field @@ -197,6 +198,15 @@ pub enum OutputEncoding { LanceJson, } +impl DeepSizeOf for OutputEncoding { + fn deep_size_of_children(&self, _context: &mut crate::deepsize::Context) -> usize { + match self { + Self::Dictionary { value, .. } => std::mem::size_of::() + value.deep_size_of(), + _ => 0, + } + } +} + impl OutputEncoding { /// The layout of an Arrow type that holds representation-only values, if /// it has one. diff --git a/rust/lance-encoding/src/data.rs b/rust/lance-encoding/src/data.rs index 84ea0300f17..2cdd4717a39 100644 --- a/rust/lance-encoding/src/data.rs +++ b/rust/lance-encoding/src/data.rs @@ -1154,6 +1154,95 @@ impl DictionaryDataBlock { } } +/// A stored layout could not be converted to the layout a reader requested, +/// for example because the values exceed the 32-bit offsets of `Utf8`. +/// +/// Returned as the source of an [`Error::InvalidInput`] so readers that know +/// the field can report it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LayoutConversionError { + /// The layout the data is stored in. + pub stored: DataType, + /// The layout the reader requested. + pub requested: DataType, + /// The bytes of values that do not fit, when known. + pub value_bytes: Option, + /// Why the conversion failed. + pub reason: String, +} + +impl std::fmt::Display for LayoutConversionError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "cannot read {} values as {}", + self.stored, self.requested + )?; + if let Some(value_bytes) = self.value_bytes { + write!(f, " ({value_bytes} bytes of values)")?; + } + write!(f, ": {}", self.reason) + } +} + +impl std::error::Error for LayoutConversionError {} + +fn dictionary_key_type(bits: u64, signed: bool) -> Option { + use DataType::*; + Some(match (bits, signed) { + (8, true) => Int8, + (16, true) => Int16, + (32, true) => Int32, + (64, true) => Int64, + (8, false) => UInt8, + (16, false) => UInt16, + (32, false) => UInt32, + (64, false) => UInt64, + _ => return None, + }) +} + +/// Convert values decoded in their stored layout to the requested layout of +/// the same values, failing instead of truncating when they do not fit. +fn convert_layout(stored: ArrayData, requested: &DataType) -> Result { + convert_layout_within(stored, requested, i32::MAX as u64) +} + +/// [`convert_layout`], where `max_offset` is the largest byte offset a layout +/// with 32-bit offsets holds. +fn convert_layout_within( + stored: ArrayData, + requested: &DataType, + max_offset: u64, +) -> Result { + let array = arrow_array::make_array(stored); + let failure = |value_bytes: Option, reason: String| { + Error::invalid_input_source(Box::new(LayoutConversionError { + stored: array.data_type().clone(), + requested: requested.clone(), + value_bytes, + reason, + })) + }; + if matches!(requested, DataType::Utf8 | DataType::Binary) { + let span = |offsets: &[i64]| (offsets[offsets.len() - 1] - offsets[0]) as u64; + let value_bytes = match array.data_type() { + DataType::LargeUtf8 => Some(span(array.as_string::().offsets())), + DataType::LargeBinary => Some(span(array.as_binary::().offsets())), + _ => None, + }; + if let Some(value_bytes) = value_bytes.filter(|bytes| *bytes > max_offset) { + return Err(failure( + Some(value_bytes), + "the values exceed 32-bit offsets".to_string(), + )); + } + } + arrow_cast::cast(&array, requested) + .map(|converted| converted.to_data()) + .map_err(|err| failure(None, err.to_string())) +} + /// A DataBlock is a collection of buffers that represents an "array" of data in very generic terms /// /// The output of each decoder is a DataBlock. Decoders can be chained together to transform @@ -1192,6 +1281,12 @@ impl DataBlock { } fn into_arrow_impl(self, data_type: DataType, validate: bool) -> Result { + if let Some(stored_type) = self.stored_layout(&data_type) + && stored_type != data_type + { + let stored = self.into_arrow_impl(stored_type, validate)?; + return convert_layout(stored, &data_type); + } match self { Self::Empty() => Ok(new_empty_array(&data_type).to_data()), Self::Constant(inner) => inner.into_arrow(data_type, validate), @@ -1208,6 +1303,54 @@ impl DataBlock { } } + /// The Arrow type this block stores when `requested` names another layout + /// of the same values. + /// + /// Each block describes its own physical layout (offset width, value width, + /// dictionary key width), so a reader may ask for any layout of the stored + /// values: strings or bytes with 32- or 64-bit offsets, views, or a + /// dictionary of them, and decimals of either width. `None` means the + /// block is decoded as `requested` directly. + fn stored_layout(&self, requested: &DataType) -> Option { + use DataType::*; + match (self, requested) { + (Self::Nullable(inner), _) => inner.data.stored_layout(requested), + (Self::VariableWidth(block), Utf8 | LargeUtf8 | Utf8View) => { + Some(if block.bits_per_offset == 64 { + LargeUtf8 + } else { + Utf8 + }) + } + (Self::VariableWidth(block), Binary | LargeBinary | BinaryView) => { + Some(if block.bits_per_offset == 64 { + LargeBinary + } else { + Binary + }) + } + (Self::VariableWidth(_), Dictionary(_, value)) => self.stored_layout(value), + ( + Self::FixedWidth(block), + Decimal128(precision, scale) | Decimal256(precision, scale), + ) => match block.bits_per_value { + 128 => Some(Decimal128(*precision, *scale)), + 256 => Some(Decimal256(*precision, *scale)), + _ => None, + }, + (Self::Dictionary(block), Dictionary(key, value)) => { + let key = + dictionary_key_type(block.indices.bits_per_value, key.is_signed_integer())?; + let value = block + .dictionary + .stored_layout(value) + .unwrap_or_else(|| value.as_ref().clone()); + Some(Dictionary(Box::new(key), Box::new(value))) + } + _ => None, + } + } + /// Convert the data block into a collection of buffers for serialization /// /// The order matters and will be used to reconstruct the data block at read time. @@ -2142,6 +2285,27 @@ mod tests { FixedWidthDataBlock, VariableWidthBlock, }; + /// Values that do not fit the requested layout fail the read with their + /// size instead of being truncated. + #[test] + fn test_convert_layout_overflow() { + let values = arrow_array::LargeStringArray::from(vec!["abc", "de"]).to_data(); + let err = super::convert_layout_within(values.clone(), &DataType::Utf8, 4).unwrap_err(); + let Error::InvalidInput { source, .. } = &err else { + panic!("{err}"); + }; + let conversion = source + .downcast_ref::() + .unwrap(); + assert_eq!(conversion.value_bytes, Some(5)); + assert_eq!(conversion.requested, DataType::Utf8); + let converted = super::convert_layout_within(values, &DataType::Utf8, 5).unwrap(); + assert_eq!( + arrow_array::make_array(converted).as_ref(), + &arrow_array::StringArray::from(vec!["abc", "de"]) as &dyn Array + ); + } + use arrow_array::Array; #[test] @@ -2847,33 +3011,46 @@ mod tests { ); } - #[test] - fn dictionary_rejects_indices_wider_than_declared_key_type() { + /// Indices stored wider than the requested key type convert when every + /// index fits, and fail the read instead of wrapping when one does not. + #[rstest] + #[case::fits(vec![0_u32, 1, 0], true)] + #[case::out_of_range(vec![0_u32, 1, 128], false)] + fn dictionary_indices_convert_to_requested_key_type( + #[case] indices: Vec, + #[case] fits: bool, + ) { + let values: Vec> = (0..=128) + .map(|i| match i { + 0 => Some("zero"), + 1 => Some("one"), + _ => None, + }) + .collect(); let dictionary = DataBlock::Dictionary(DictionaryDataBlock { indices: FixedWidthDataBlock { - data: LanceBuffer::reinterpret_vec(vec![0_u32, 1, 128]), + data: LanceBuffer::reinterpret_vec(indices), bits_per_value: 32, num_values: 3, block_info: BlockInfo::new(), }, - dictionary: Box::new(DataBlock::from_array(StringArray::from(vec![ - Some("zero"), - Some("one"), - None, - ]))), + dictionary: Box::new(DataBlock::from_array(StringArray::from(values))), }); let data_type = DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)); - let error = dictionary - .into_arrow(data_type, false) - .expect_err("mismatched dictionary index widths must be rejected"); - - assert!(matches!(error, Error::CorruptFile { .. })); - assert!( - error.to_string().contains( - "dictionary indices use 32 bits but the declared Int8 key type uses 8 bits" - ) - ); + let result = dictionary.into_arrow(data_type.clone(), false); + if fits { + let array = make_array(result.unwrap()); + assert_eq!(array.data_type(), &data_type); + let expected: ArrayRef = Arc::new(StringArray::from(vec!["zero", "one", "zero"])); + assert_eq!( + arrow_cast::cast(&array, &DataType::Utf8).unwrap().as_ref(), + expected.as_ref() + ); + } else { + let error = result.expect_err("an index outside the key type must fail"); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + } } #[rstest] diff --git a/rust/lance-encoding/src/decoder.rs b/rust/lance-encoding/src/decoder.rs index a3732f1e47f..7231c89a5ec 100644 --- a/rust/lance-encoding/src/decoder.rs +++ b/rust/lance-encoding/src/decoder.rs @@ -716,13 +716,16 @@ impl CoreFieldDecoderStrategy { match data_type { // DataType::is_primitive doesn't consider these primitive but we do DataType::Dictionary(_, value_type) => Self::is_structural_primitive(value_type), + // Views are decoded from the offset layout they are stored as. DataType::Boolean | DataType::Null | DataType::FixedSizeBinary(_) | DataType::Binary | DataType::LargeBinary + | DataType::BinaryView | DataType::Utf8 - | DataType::LargeUtf8 => true, + | DataType::LargeUtf8 + | DataType::Utf8View => true, DataType::FixedSizeList(inner, _) => { Self::is_structural_primitive(inner.data_type()) } diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 0cab0801e4b..4c36a335d42 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -5285,6 +5285,17 @@ impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler { } } +/// Name the field in an error about converting its stored layout to the +/// requested one; other errors pass through unchanged. +fn name_layout_error(err: Error, field_name: &str) -> Error { + if let Error::InvalidInput { source, .. } = &err + && let Some(conversion) = source.downcast_ref::() + { + return Error::invalid_input(format!("field '{field_name}': {conversion}")); + } + err +} + /// Takes the output from several pages decoders and /// concatenates them. #[derive(Debug)] @@ -5292,6 +5303,8 @@ pub struct StructuralCompositeDecodeArrayTask { tasks: Vec>, should_validate: bool, data_type: DataType, + /// Names the field in errors about the requested layout. + field_name: String, } impl StructuralCompositeDecodeArrayTask { @@ -5342,7 +5355,8 @@ impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask { let array = make_array( decoded .data - .into_arrow(self.data_type.clone(), self.should_validate)?, + .into_arrow(self.data_type.clone(), self.should_validate) + .map_err(|err| name_layout_error(err, &self.field_name))?, ); arrays.push(array); @@ -5468,6 +5482,7 @@ impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder { tasks, should_validate: self.should_validate, data_type: self.field.data_type().clone(), + field_name: self.field.name().clone(), })) } diff --git a/rust/lance-encoding/src/encodings/logical/primitive/constant.rs b/rust/lance-encoding/src/encodings/logical/primitive/constant.rs index 3ac011ceb49..59039e468fb 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/constant.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/constant.rs @@ -205,6 +205,42 @@ impl ConstantPageScheduler { } } +/// The layout a constant page's scalar was written in, for a reader that +/// requested `requested`. +/// +/// A column's files may store different layouts of the same values, and a +/// scalar records only its buffers, whose lengths tell the layouts apart: one +/// value has an offsets buffer of 8 bytes with 32-bit offsets and 16 bytes with +/// 64-bit offsets, and a decimal value takes 16 or 32 bytes. +fn stored_scalar_type(requested: &DataType, buffer_lengths: &[usize]) -> DataType { + use DataType::*; + let wide_offsets = buffer_lengths.first() == Some(&16); + match requested { + Utf8 | LargeUtf8 | Utf8View => { + if wide_offsets { + LargeUtf8 + } else { + Utf8 + } + } + Binary | LargeBinary | BinaryView => { + if wide_offsets { + LargeBinary + } else { + Binary + } + } + Dictionary(_, value) => stored_scalar_type(value, buffer_lengths), + Decimal128(precision, scale) | Decimal256(precision, scale) => { + match buffer_lengths.first() { + Some(32) => Decimal256(*precision, *scale), + _ => Decimal128(*precision, *scale), + } + } + _ => requested.clone(), + } +} + impl crate::encodings::logical::primitive::StructuralPageScheduler for ConstantPageScheduler { fn init_layout(&self) -> Result { // Order must match `init_from_buffers`' consumption: scalar, rep, def. @@ -229,11 +265,18 @@ impl crate::encodings::logical::primitive::StructuralPageScheduler for ConstantP let scalar = match (scalar_source, scalar_buffer) { (ScalarSource::Inline(inline), None) => { - lance_arrow::scalar::decode_scalar_from_inline_value(&data_type, &inline)? + let stored_type = stored_scalar_type(&data_type, &[inline.len()]); + lance_arrow::scalar::decode_scalar_from_inline_value(&stored_type, &inline)? } (ScalarSource::ValueBuffer(_), Some(bytes)) => { let buf = LanceBuffer::from_bytes(bytes, 1); - lance_arrow::scalar::decode_scalar_from_value_buffer(&data_type, buf.as_ref())? + let buffer_lengths = + lance_arrow::scalar::scalar_value_buffer_lengths(buf.as_ref())?; + let stored_type = stored_scalar_type(&data_type, &buffer_lengths); + lance_arrow::scalar::decode_scalar_from_value_buffer( + &stored_type, + buf.as_ref(), + )? } _ => { return Err(Error::internal( @@ -241,6 +284,11 @@ impl crate::encodings::logical::primitive::StructuralPageScheduler for ConstantP )); } }; + let scalar = if scalar.data_type() == &data_type { + scalar + } else { + arrow_cast::cast(&scalar, &data_type)? + }; let rep = rep_buffer.map(|rep| { let rep = LanceBuffer::from_bytes(rep, 2); diff --git a/rust/lance-file/src/datatypes.rs b/rust/lance-file/src/datatypes.rs index 3d84e99267c..10c0cb8e05b 100644 --- a/rust/lance-file/src/datatypes.rs +++ b/rust/lance-file/src/datatypes.rs @@ -2,7 +2,9 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use lance_arrow::ARROW_EXT_NAME_KEY; -use lance_core::datatypes::{Dictionary, Encoding, Field, LogicalType, Schema}; +use lance_core::datatypes::{ + Dictionary, Encoding, Field, LogicalType, OUTPUT_ENCODING_META_KEY, Schema, +}; use lance_core::{Error, Result}; use std::collections::HashMap; @@ -51,17 +53,21 @@ impl From<&pb::Field> for Field { } else { None }, + output_encoding: None, } } } impl From<&Field> for pb::Field { fn from(field: &Field) -> Self { - let pb_metadata = field + let mut pb_metadata: HashMap> = field .metadata .iter() .map(|(key, value)| (key.clone(), value.clone().into_bytes())) .collect(); + if let Some(entry) = field.output_encoding_entry() { + pb_metadata.insert(OUTPUT_ENCODING_META_KEY.to_string(), entry.into_bytes()); + } Self { id: field.id, parent_id: field.parent_id, diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 65aa19939a7..c5d5dac253a 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -4732,3 +4732,170 @@ mod tests { assert_eq!(validate(empty_struct, true, &[0], vec![9]).unwrap(), 9); } } + +/// A reader may request any layout of the values a column stores: each page +/// describes its own offset width, value width, and dictionary key width, so +/// files of one column can store different layouts. +#[cfg(test)] +mod layout_request_tests { + use std::sync::Arc; + + use arrow_array::{ + ArrayRef, BinaryArray, Decimal128Array, RecordBatch, RecordBatchIterator, StringArray, + }; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + use futures::TryStreamExt; + use lance_core::datatypes::{Field, OutputEncoding, Schema}; + use lance_encoding::decoder::{DecoderPlugins, FilterExpression}; + use lance_io::utils::CachedFileSize; + use rstest::rstest; + + use crate::reader::{FileReader, FileReaderOptions, ReaderProjection}; + use crate::testing::{FsFixture, test_cache, write_lance_file}; + use crate::version::ConcreteFileVersion; + use crate::writer::FileWriterOptions; + + fn dictionary(key: DataType, value: DataType) -> DataType { + DataType::Dictionary(Box::new(key), Box::new(value)) + } + + /// Stored arrays, each with the layouts it can be read as. + fn cases() -> Vec<(ArrayRef, Vec)> { + let strings: ArrayRef = Arc::new(StringArray::from(vec![ + Some("a"), + None, + Some("bb"), + Some("a"), + ])); + let bytes: ArrayRef = Arc::new(BinaryArray::from(vec![ + Some(b"a".as_ref()), + None, + Some(b"bb".as_ref()), + ])); + let decimals: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(1), None, Some(-12345)]) + .with_precision_and_scale(10, 2) + .unwrap(), + ); + let string_layouts = vec![ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + dictionary(DataType::Int16, DataType::Utf8), + dictionary(DataType::UInt32, DataType::LargeUtf8), + ]; + let binary_layouts = vec![ + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + dictionary(DataType::Int8, DataType::Binary), + ]; + let decimal_layouts = vec![DataType::Decimal128(10, 2), DataType::Decimal256(10, 2)]; + // Repeated values are written as constant pages, whose scalar records + // only its own buffers. + let constant_strings: ArrayRef = Arc::new(StringArray::from(vec!["a"; 4])); + let constant_decimals: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(7); 4]) + .with_precision_and_scale(10, 2) + .unwrap(), + ); + let mut cases = Vec::new(); + for (values, layouts) in [ + (strings, string_layouts.clone()), + (constant_strings, string_layouts), + (bytes, binary_layouts), + (decimals, decimal_layouts.clone()), + (constant_decimals, decimal_layouts), + ] { + for stored in layouts + .iter() + .filter(|layout| !matches!(layout, DataType::Utf8View | DataType::BinaryView)) + { + cases.push((arrow_cast::cast(&values, stored).unwrap(), layouts.clone())); + } + } + cases + } + + fn requested_field(data_type: &DataType) -> Field { + let mut field = Field::try_from(ArrowField::new("c", data_type.clone(), true)).unwrap(); + field.id = 0; + if matches!(data_type, DataType::Utf8View | DataType::BinaryView) { + field.output_encoding = OutputEncoding::of_data_type(data_type); + } + field + } + + #[rstest] + #[case::v2_1(ConcreteFileVersion::V2_1)] + #[case::v2_2(ConcreteFileVersion::V2_2)] + #[case::v2_3(ConcreteFileVersion::V2_3)] + #[tokio::test] + async fn test_read_any_layout_of_stored_values(#[case] version: ConcreteFileVersion) { + for (stored, requested_layouts) in cases() { + let fs = FsFixture::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "c", + stored.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![stored.clone()]).unwrap(); + write_lance_file( + RecordBatchIterator::new(vec![Ok(batch)], schema), + &fs, + version, + FileWriterOptions::default(), + ) + .await; + for requested in requested_layouts { + let file = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let reader = FileReader::try_open( + file, + None, + Arc::::default(), + &test_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + let projection = ReaderProjection { + column_indices: vec![0], + schema: Arc::new(Schema { + fields: vec![requested_field(&requested)], + metadata: Default::default(), + }), + }; + let batches = reader + .read_stream_projected( + lance_io::ReadBatchParams::RangeFull, + 1024, + 16, + projection, + FilterExpression::no_filter(), + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap_or_else(|err| panic!("{} as {requested}: {err}", stored.data_type())); + let read = batches[0].column(0); + assert_eq!( + read.data_type(), + &requested, + "{} as {requested}", + stored.data_type() + ); + assert_eq!( + arrow_cast::cast(read, stored.data_type()).unwrap().as_ref(), + stored.as_ref(), + "{} as {requested}", + stored.data_type() + ); + } + } + } +} diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 3c67753910d..21cd83c6fc9 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -5,6 +5,7 @@ use crate::format::Manifest; use lance_core::{Error, Result}; +use lance_file::version::ConcreteFileVersion; /// Fragments may contain deletion files, which record the tombstones of /// soft-deleted rows. @@ -64,8 +65,18 @@ pub const FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS: u64 = 1 << 9; /// preserves them during maintenance. Legacy-only FRI does not set this bit. /// Bit 9 is taken by the stable-row-id FRI compatibility flag. pub const FLAG_FRAGMENT_REUSE_INDEX: u64 = 1 << 10; +/// The table schema follows the semantic type contract: `logical_type` names a +/// semantic type, `lance-schema:output-encoding` selects the Arrow layout of +/// reads, and data files of one column may use different physical layouts. +/// +/// A reader without this bit would return other Arrow types than the table +/// specifies and cannot parse canonical names such as `decimal:10:2`; a writer +/// without it would reject appends the contract accepts and write legacy +/// aliases. Set only when a table is created with data storage version 2.3 or +/// later, and kept in every later version. +pub const FLAG_SEMANTIC_TYPES: u64 = 1 << 11; /// The first bit that is unknown as a feature flag -pub const FLAG_UNKNOWN: u64 = 1 << 11; +pub const FLAG_UNKNOWN: u64 = 1 << 12; const _: () = assert!(FLAG_COVERED_INDEX_METADATA < FLAG_UNKNOWN); // The fence needs a bit the current released build already refuses, which means @@ -77,8 +88,28 @@ const _: () = assert!(FLAG_MIXED_DATA_FILE_VERSIONS < FLAG_UNKNOWN); const _: () = assert!(FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS >= 1 << 8); const _: () = assert!(FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS < FLAG_UNKNOWN); const _: () = assert!(FLAG_FRAGMENT_REUSE_INDEX < FLAG_UNKNOWN); +// Builds released before the semantic type contract treat bit 11 and above as +// unknown, so they refuse flagged tables. +const _: () = assert!(FLAG_SEMANTIC_TYPES >= 1 << 11); +const _: () = assert!(FLAG_SEMANTIC_TYPES < FLAG_UNKNOWN); + +/// Capabilities whose reader and writer bits are set together and, once set, +/// stay set in every later version, including restores. +pub(crate) const STICKY_PAIRED_FLAGS: u64 = FLAG_MIXED_DATA_FILE_VERSIONS | FLAG_SEMANTIC_TYPES; -pub(crate) const STICKY_PAIRED_FLAGS: u64 = FLAG_MIXED_DATA_FILE_VERSIONS; +/// Whether a table created with data storage `version` follows the semantic +/// type contract and sets [`FLAG_SEMANTIC_TYPES`]. +/// +/// Existing tables never switch: the flag is decided once, at creation. +pub fn creates_semantic_types(version: ConcreteFileVersion) -> bool { + match version { + ConcreteFileVersion::V1 + | ConcreteFileVersion::V2_0 + | ConcreteFileVersion::V2_1 + | ConcreteFileVersion::V2_2 => false, + ConcreteFileVersion::V2_3 => true, + } +} /// Environment variable that opts a release build into reading and writing data /// overlay files before the feature is generally released. @@ -272,14 +303,21 @@ pub fn has_deprecated_v2_feature_flag(writer_flags: u64) -> bool { /// commit path refuses to *produce* this, so seeing it on read means the /// manifest was written by something that did not. pub fn validate_paired_feature_flags(manifest: &Manifest) -> Result<()> { - let reader = manifest.reader_feature_flags & FLAG_MIXED_DATA_FILE_VERSIONS != 0; - let writer = manifest.writer_feature_flags & FLAG_MIXED_DATA_FILE_VERSIONS != 0; - if reader != writer { - return Err(Error::corrupt_file_named( - "manifest", - "Manifest has only one of the mixed data-file-version reader and writer feature bits set, \ - so its semantics are undefined", - )); + for (flag, capability) in [ + (FLAG_MIXED_DATA_FILE_VERSIONS, "mixed data-file-version"), + (FLAG_SEMANTIC_TYPES, "semantic type"), + ] { + let reader = manifest.reader_feature_flags & flag != 0; + let writer = manifest.writer_feature_flags & flag != 0; + if reader != writer { + return Err(Error::corrupt_file_named( + "manifest", + format!( + "Manifest has only one of the {capability} reader and writer feature bits set, \ + so its semantics are undefined" + ), + )); + } } Ok(()) } @@ -614,6 +652,46 @@ mod tests { ) } + /// This build reads and writes tables under the semantic type contract, + /// while a build released before it, whose unknown boundary is bit 11, + /// refuses them. + #[test] + fn semantic_types_flag_fences_older_builds() { + assert!(can_read_dataset(FLAG_SEMANTIC_TYPES)); + assert!(can_write_dataset(FLAG_SEMANTIC_TYPES)); + let supported_before_contract = (1 << 11) - 1; + assert_ne!(FLAG_SEMANTIC_TYPES & !supported_before_contract, 0); + } + + #[rstest::rstest] + #[case::v1(ConcreteFileVersion::V1, false)] + #[case::v2_0(ConcreteFileVersion::V2_0, false)] + #[case::v2_1(ConcreteFileVersion::V2_1, false)] + #[case::v2_2(ConcreteFileVersion::V2_2, false)] + #[case::v2_3(ConcreteFileVersion::V2_3, true)] + fn semantic_types_start_at_v2_3(#[case] version: ConcreteFileVersion, #[case] expected: bool) { + assert_eq!(creates_semantic_types(version), expected); + } + + #[test] + fn semantic_types_flag_is_sticky_and_paired() { + let mut manifest = empty_manifest(); + manifest.reader_feature_flags = FLAG_SEMANTIC_TYPES; + manifest.writer_feature_flags = FLAG_SEMANTIC_TYPES; + apply_feature_flags(&mut manifest, false, false).unwrap(); + assert_ne!(manifest.reader_feature_flags & FLAG_SEMANTIC_TYPES, 0); + assert_ne!(manifest.writer_feature_flags & FLAG_SEMANTIC_TYPES, 0); + + let mut derived = empty_manifest(); + inherit_sticky_feature_flags(&mut derived, &manifest).unwrap(); + assert_ne!(derived.reader_feature_flags & FLAG_SEMANTIC_TYPES, 0); + assert_ne!(derived.writer_feature_flags & FLAG_SEMANTIC_TYPES, 0); + + manifest.writer_feature_flags = 0; + let err = validate_paired_feature_flags(&manifest).unwrap_err(); + assert!(err.to_string().contains("semantic type"), "{err}"); + } + #[test] fn mixed_capability_is_below_the_unknown_boundary() { assert!(can_read_dataset(FLAG_COVERED_INDEX_METADATA)); diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index b8197c84271..b1a219a46b3 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -20,7 +20,9 @@ use std::sync::Arc; use super::{Fragment, InlineRowIds, RowIdMeta}; use crate::feature_flags::{FLAG_COVERED_INDEX_METADATA, STICKY_PAIRED_FLAGS}; -use crate::feature_flags::{FLAG_STABLE_ROW_IDS, has_deprecated_v2_feature_flag}; +use crate::feature_flags::{ + FLAG_SEMANTIC_TYPES, FLAG_STABLE_ROW_IDS, has_deprecated_v2_feature_flag, +}; use crate::format::fragment::DataFileFieldInterner; use crate::format::pb; use lance_core::cache::LanceCache; @@ -544,12 +546,23 @@ impl Manifest { self.reader_feature_flags & FLAG_STABLE_ROW_IDS != 0 } + /// Whether the table schema follows the semantic type contract + /// ([`FLAG_SEMANTIC_TYPES`]). + pub fn uses_semantic_types(&self) -> bool { + self.reader_feature_flags & FLAG_SEMANTIC_TYPES != 0 + } + /// How this table compares field types for schema compatibility. /// - /// Every table currently names exactly one Arrow type per `logical_type`, - /// so types are compared exactly. + /// Legacy tables name exactly one Arrow type per `logical_type` and compare + /// types exactly; tables under the semantic type contract compare semantic + /// types. pub fn type_comparison(&self) -> TypeComparison { - TypeComparison::Exact + if self.uses_semantic_types() { + TypeComparison::Semantic + } else { + TypeComparison::Exact + } } /// Creates a serialized copy of the manifest, suitable for IPC or temp storage @@ -1022,6 +1035,13 @@ impl TryFrom for Manifest { }; let schema = Schema::try_from(fields_with_meta)?; + // A table under the semantic type contract holds its schema in + // canonical form, with output encodings resolved from field metadata. + let schema = if p.reader_feature_flags & FLAG_SEMANTIC_TYPES != 0 { + schema.to_canonical_types()? + } else { + schema + }; Ok(Self { schema, diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index fc8dc4e9107..9773b60be5a 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -11,8 +11,9 @@ //! metadata it stamps, the validation that runs before it. use crate::feature_flags::{ - FLAG_COVERED_INDEX_METADATA, FLAG_STABLE_ROW_IDS, apply_feature_flags, - ensure_can_read_manifest, ensure_can_write_manifest, inherit_sticky_feature_flags, + FLAG_COVERED_INDEX_METADATA, FLAG_SEMANTIC_TYPES, FLAG_STABLE_ROW_IDS, apply_feature_flags, + creates_semantic_types, ensure_can_read_manifest, ensure_can_write_manifest, + inherit_sticky_feature_flags, }; use crate::format::overlay::{OverlayCoverage, TOMBSTONE_FIELD_ID}; use crate::format::{ @@ -41,7 +42,7 @@ use crate::transaction::{ }; use lance_core::datatypes::{ LANCE_UNENFORCED_CLUSTERING_KEY_POSITION, LANCE_UNENFORCED_PRIMARY_KEY, - LANCE_UNENFORCED_PRIMARY_KEY_POSITION, + LANCE_UNENFORCED_PRIMARY_KEY_POSITION, OUTPUT_ENCODING_META_KEY, }; use lance_core::utils::parse::str_is_truthy; use lance_core::{Error, Result}; @@ -197,6 +198,11 @@ impl Transaction { ))); } inherit_sticky_feature_flags(&mut manifest, current_manifest)?; + if manifest.uses_semantic_types() { + // A version from before the table adopted the contract holds legacy + // aliases, which read as their canonical types and output encodings. + manifest.schema = manifest.schema.to_canonical_types()?; + } Ok((manifest, indices)) } @@ -1462,6 +1468,15 @@ impl Transaction { manifest.tag.clone_from(&self.tag); + // The contract is decided once, when the table is created; the flag is + // sticky, so apply_feature_flags and inheritance keep it from then on. + if current_manifest.is_none() + && creates_semantic_types(manifest.data_storage_format.lance_file_format()) + { + manifest.reader_feature_flags |= FLAG_SEMANTIC_TYPES; + manifest.writer_feature_flags |= FLAG_SEMANTIC_TYPES; + } + if config.auto_set_feature_flags { // Internal operations (e.g. CreateIndex) build with the default config, // which has use_stable_row_ids = false. Without inheriting from the previous @@ -1566,8 +1581,20 @@ impl Transaction { .iter() .any(|entry| entry.key == LANCE_UNENFORCED_CLUSTERING_KEY_POSITION) }); + let semantic_types = manifest.uses_semantic_types(); for (field_id, field_metadata_update) in field_metadata_updates { if let Some(field) = manifest.schema.field_by_id_mut(*field_id) { + if semantic_types { + // Updates see the output encoding as the metadata + // entry it is recorded as; the schema is converted + // back to canonical form once all updates apply. + if let Some(entry) = field.output_encoding_entry() { + field + .metadata + .insert(OUTPUT_ENCODING_META_KEY.to_string(), entry); + } + field.output_encoding = None; + } apply_update_map(&mut field.metadata, field_metadata_update); // Also set unenforced primary key based on updated field metadata. field.unenforced_primary_key_position = field @@ -1690,6 +1717,17 @@ impl Transaction { manifest.next_row_id = next_row_id; } + if manifest.uses_semantic_types() { + // Output encodings are checked when set; entries carried over + // unchanged were written by other builds and are left alone. + manifest.schema.check_output_encoding_entries( + current_manifest + .filter(|current| current.uses_semantic_types()) + .map(|current| ¤t.schema), + )?; + manifest.schema = manifest.schema.to_canonical_types()?; + } + Ok((manifest, final_indices)) } diff --git a/rust/lance/src/dataset/fragment/write.rs b/rust/lance/src/dataset/fragment/write.rs index 8c153166127..38552ebccff 100644 --- a/rust/lance/src/dataset/fragment/write.rs +++ b/rust/lance/src/dataset/fragment/write.rs @@ -164,12 +164,13 @@ impl<'a> FragmentCreateBuilder<'a> { let mut fragment = Fragment::new(id); let full_path = base_path.clone().join(DATA_DIR).join(filename.clone()); let obj_writer = object_store.create(&full_path).await?; - let (writer, data_file) = create_writer(obj_writer, schema, filename)?; + let file_schema = schema.to_data_file_schema()?; + let (writer, data_file) = create_writer(obj_writer, file_schema.clone(), filename)?; fragment.files.push(data_file.clone()); progress.begin(&fragment).await?; - let mut writer = V2WriterAdapter::new(writer, Some(data_file), None); + let mut writer = V2WriterAdapter::new(writer, Some(data_file), None, &file_schema); let break_limit = (128 * 1024).min(params.max_rows_per_file); let mut broken_stream = break_stream(stream, break_limit); diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index 9632d338035..adb4a2bd81f 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -681,6 +681,13 @@ pub(super) async fn can_use_binary_copy_current( ); return Ok(false); } + // Under the semantic type contract, data files of one column may hold + // different layouts, and copied pages keep theirs. + let data_file_schema = dataset + .manifest + .uses_semantic_types() + .then(|| dataset.schema().to_data_file_schema()) + .transpose()?; for fragment in fragments { // Binary copy only reads base files; overlays must be materialized by the scanner. if !fragment.overlays.is_empty() { @@ -732,6 +739,21 @@ pub(super) async fn can_use_binary_copy_current( ); return Ok(false); } + if let Some(data_file_schema) = &data_file_schema { + let differs = data_file_schema.fields_pre_order().any(|field| { + file_meta + .file_schema + .field_by_id(field.id) + .is_some_and(|file_field| file_field.logical_type != field.logical_type) + }); + if differs { + log::debug!( + "Binary copy disabled: data file {} stores a column in another layout", + data_file.path + ); + return Ok(false); + } + } } } diff --git a/rust/lance/src/dataset/optimize/binary_copy.rs b/rust/lance/src/dataset/optimize/binary_copy.rs index c76e0ea300f..a151682bd7d 100644 --- a/rust/lance/src/dataset/optimize/binary_copy.rs +++ b/rust/lance/src/dataset/optimize/binary_copy.rs @@ -141,7 +141,9 @@ pub async fn rewrite_files_binary_copy( // - Writes a new footer (schema descriptor, column metadata, offset tables, version) // - Optionally carries forward stable row ids and persists them inline in fragment metadata // Merge small Lance files into larger ones by page-level binary copy. - let schema = dataset.schema().clone(); + // The copied pages keep their layouts, which eligibility requires to be + // the ones the table schema's data file form records. + let schema = dataset.schema().to_data_file_schema()?; let column_count = schema .fields .iter() diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index f72c6f1b8c6..21683efaa05 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1378,7 +1378,15 @@ async fn test_write_manifest( manifest.data_storage_format.version.to_manifest_string(), "stable" | "next" )); - assert_eq!(manifest.reader_feature_flags, 0); + // A new table follows the semantic type contract from data storage 2.3 on. + let creation_flags = if feature_flags::creates_semantic_types( + manifest.data_storage_format.lance_file_format(), + ) { + feature_flags::FLAG_SEMANTIC_TYPES + } else { + 0 + }; + assert_eq!(manifest.reader_feature_flags, creation_flags); // Create one with deletions dataset.delete("i < 10").await.unwrap(); @@ -1399,11 +1407,11 @@ async fn test_write_manifest( .unwrap(); assert_eq!( manifest.writer_feature_flags, - feature_flags::FLAG_DELETION_FILES + feature_flags::FLAG_DELETION_FILES | creation_flags ); assert_eq!( manifest.reader_feature_flags, - feature_flags::FLAG_DELETION_FILES + feature_flags::FLAG_DELETION_FILES | creation_flags ); // Write with custom manifest @@ -4170,3 +4178,591 @@ async fn a_legacy_nullable_primary_key_can_be_repaired_under_mem_wal() { .expect("removing the offending rows must not be blocked under MemWAL"); assert_eq!(dataset.count_rows(None).await.unwrap(), 1); } + +fn string_array(data_type: &DataType, values: &[Option<&str>]) -> ArrayRef { + let utf8: ArrayRef = Arc::new(StringArray::from(values.to_vec())); + arrow_cast::cast(&utf8, data_type).unwrap() +} + +/// Write one column `s` of `array` to `uri` and return the dataset. +async fn write_single_column( + uri: &str, + array: ArrayRef, + mode: WriteMode, + version: Option, +) -> Result { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "s", + array.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + uri, + Some(WriteParams { + mode, + data_storage_version: version, + ..Default::default() + }), + ) + .await +} + +/// The Arrow type each fragment's data file records for `column`. +async fn data_file_layouts(dataset: &Dataset, column: &str) -> Vec { + use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; + use lance_io::utils::CachedFileSize; + + let scheduler = ScanScheduler::new( + dataset.object_store.clone(), + SchedulerConfig::default_for_testing(), + ); + let mut layouts = Vec::new(); + for fragment in dataset.manifest.fragments.iter() { + let data_file = &fragment.files[0]; + let path = dataset + .data_file_dir(data_file) + .unwrap() + .join(data_file.path.as_str()); + let file = scheduler + .open_file(&path, &CachedFileSize::unknown()) + .await + .unwrap(); + let metadata = lance_file::reader::FileReader::read_all_metadata(&file) + .await + .unwrap(); + layouts.push(metadata.file_schema.field(column).unwrap().data_type()); + } + layouts +} + +/// On a table under the semantic type contract, a `string` column accepts +/// every string layout. Each data file keeps the layout it was written in, and +/// reads return the layout the column was created with. +#[rstest] +#[case::utf8(DataType::Utf8, None)] +#[case::large_utf8(DataType::LargeUtf8, Some("large_utf8"))] +#[case::utf8_view(DataType::Utf8View, Some("utf8_view"))] +#[case::dictionary( + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + Some("dictionary:int16:utf8") +)] +#[tokio::test] +async fn test_semantic_string_column_accepts_every_layout( + #[case] created: DataType, + #[case] output_encoding: Option<&str>, +) { + let dir = TempStrDir::default(); + let uri = dir.as_str(); + let appended = [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + ]; + let mut expected = vec![Some("a"), None, Some("a")]; + write_single_column( + uri, + string_array(&created, &expected), + WriteMode::Create, + Some(LanceFileVersion::V2_3), + ) + .await + .unwrap(); + for (index, layout) in appended.iter().enumerate() { + let value = format!("v{index}"); + let values = [Some(value.as_str()), None]; + write_single_column(uri, string_array(layout, &values), WriteMode::Append, None) + .await + .unwrap(); + expected.push(Some(Box::leak(value.into_boxed_str()))); + expected.push(None); + } + + let dataset = Dataset::open(uri).await.unwrap(); + assert_ne!( + dataset.manifest.reader_feature_flags & feature_flags::FLAG_SEMANTIC_TYPES, + 0 + ); + assert_ne!( + dataset.manifest.writer_feature_flags & feature_flags::FLAG_SEMANTIC_TYPES, + 0 + ); + let field = dataset.schema().field("s").unwrap(); + assert_eq!(field.logical_type.to_string(), "string"); + assert_eq!(field.output_encoding_entry().as_deref(), output_encoding); + + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(batch.column(0).data_type(), &created); + assert_eq!( + batch.column(0).as_ref(), + string_array(&created, &expected).as_ref() + ); + + // A view is stored as the offset layout that holds every value. + let stored = |layout: &DataType| match layout { + DataType::Utf8View => DataType::LargeUtf8, + layout => layout.clone(), + }; + let mut expected_layouts = vec![stored(&created)]; + expected_layouts.extend(appended.iter().map(stored)); + assert_eq!(data_file_layouts(&dataset, "s").await, expected_layouts); +} + +/// Compaction and scalar indices read a column whose data files hold different +/// layouts, and see the same values in the table's output layout. +#[rstest] +#[case::reencode(crate::dataset::optimize::CompactionMode::Reencode)] +#[case::try_binary_copy(crate::dataset::optimize::CompactionMode::TryBinaryCopy)] +#[tokio::test] +async fn test_semantic_mixed_layouts_compact_and_index( + #[case] compaction_mode: crate::dataset::optimize::CompactionMode, +) { + use crate::dataset::optimize::{CompactionOptions, compact_files}; + + let dir = TempStrDir::default(); + let uri = dir.as_str(); + let writes = [ + (DataType::Utf8, [Some("a"), Some("b")]), + (DataType::LargeUtf8, [Some("c"), Some("a")]), + (DataType::Utf8View, [Some("a"), None]), + ( + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + [Some("e"), Some("a")], + ), + ]; + for (index, (layout, values)) in writes.iter().enumerate() { + let mode = if index == 0 { + WriteMode::Create + } else { + WriteMode::Append + }; + write_single_column( + uri, + string_array(layout, values), + mode, + Some(LanceFileVersion::V2_3), + ) + .await + .unwrap(); + } + let expected = string_array( + &DataType::Utf8, + &writes + .iter() + .flat_map(|(_, values)| values.iter().copied()) + .collect::>(), + ); + + let mut dataset = Dataset::open(uri).await.unwrap(); + dataset + .create_index( + &["s"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + async fn count_indexed(dataset: &Dataset, predicate: &str) -> usize { + let mut scan = dataset.scan(); + scan.filter(predicate).unwrap(); + let plan = scan.explain_plan(false).await.unwrap(); + assert!(plan.contains("ScalarIndexQuery"), "{plan}"); + scan.try_into_batch().await.unwrap().num_rows() + } + assert_eq!(count_indexed(&dataset, "s = 'a'").await, 4); + assert_eq!(count_indexed(&dataset, "s = 'e'").await, 1); + + compact_files( + &mut dataset, + CompactionOptions { + compaction_mode: Some(compaction_mode), + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), 1); + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(batch.column(0).as_ref(), expected.as_ref()); + assert_eq!(count_indexed(&dataset, "s = 'a'").await, 4); + assert_eq!(count_indexed(&dataset, "s = 'c'").await, 1); +} + +/// A decimal column accepts either width that holds its precision, and +/// rejects another precision, which is a different value domain. +#[tokio::test] +async fn test_semantic_decimal_column_widths() { + use arrow_array::{Decimal128Array, Decimal256Array}; + use arrow_buffer::i256; + + let dir = TempStrDir::default(); + let uri = dir.as_str(); + let decimal128 = |precision: u8, values: Vec>| -> ArrayRef { + Arc::new( + Decimal128Array::from(values) + .with_precision_and_scale(precision, 2) + .unwrap(), + ) + }; + write_single_column( + uri, + decimal128(10, vec![Some(100), None]), + WriteMode::Create, + Some(LanceFileVersion::V2_3), + ) + .await + .unwrap(); + let wider: ArrayRef = Arc::new( + Decimal256Array::from(vec![Some(i256::from(-250))]) + .with_precision_and_scale(10, 2) + .unwrap(), + ); + write_single_column(uri, wider, WriteMode::Append, None) + .await + .unwrap(); + let err = write_single_column(uri, decimal128(12, vec![Some(1)]), WriteMode::Append, None) + .await + .unwrap_err(); + assert!( + err.to_string().contains("decimal:10:2") && err.to_string().contains("decimal:128:12:2"), + "{err}" + ); + + let dataset = Dataset::open(uri).await.unwrap(); + assert_eq!( + dataset + .schema() + .field("s") + .unwrap() + .logical_type + .to_string(), + "decimal:10:2" + ); + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!( + batch.column(0).as_ref(), + decimal128(10, vec![Some(100), None, Some(-250)]).as_ref() + ); + assert_eq!( + data_file_layouts(&dataset, "s").await, + vec![DataType::Decimal128(10, 2), DataType::Decimal256(10, 2)] + ); +} + +/// Tables on an earlier data storage version stay legacy tables: the schema +/// keeps the Arrow-mapped strings, and appends compare exact types. +#[tokio::test] +async fn test_legacy_storage_version_keeps_exact_types() { + use arrow_array::Decimal256Array; + use arrow_buffer::i256; + + let dir = TempStrDir::default(); + let uri = dir.as_str(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("s", DataType::LargeUtf8, true), + ArrowField::new( + "d", + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + true, + ), + ArrowField::new("n", DataType::Decimal256(10, 2), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + string_array(&DataType::LargeUtf8, &[Some("a")]), + string_array(schema.field(1).data_type(), &[Some("b")]), + Arc::new( + Decimal256Array::from(vec![Some(i256::from(1))]) + .with_precision_and_scale(10, 2) + .unwrap(), + ), + ], + ) + .unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch.clone())], schema.clone()), + uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!( + dataset.manifest.reader_feature_flags & feature_flags::FLAG_SEMANTIC_TYPES, + 0 + ); + let fields = pb::Manifest::from(dataset.manifest.as_ref()).fields; + let recorded = fields + .iter() + .map(|field| (field.logical_type.as_str(), field.metadata.is_empty())) + .collect::>(); + assert_eq!( + recorded, + vec![ + ("large_string", true), + ("dict:string:int16:false", true), + ("decimal:256:10:2", true), + ] + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), batch); + + let err = write_single_column( + uri, + string_array(&DataType::Utf8, &[Some("c")]), + WriteMode::Append, + None, + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("should have type large_string but type was string"), + "{err}" + ); +} + +/// A 2.3 table created before the semantic type contract has no flag, keeps +/// comparing exact types, and is never upgraded implicitly. +#[tokio::test] +async fn test_legacy_table_on_current_storage_version_is_not_upgraded() { + let dir = TempStrDir::default(); + let uri = dir.as_str(); + let dataset = write_single_column( + uri, + string_array(&DataType::Utf8, &[Some("a")]), + WriteMode::Create, + Some(LanceFileVersion::V2_3), + ) + .await + .unwrap(); + // `string` names `Utf8` in both vocabularies, so clearing the flag leaves + // the table exactly as an earlier build wrote it. + let mut manifest = dataset.manifest.as_ref().clone(); + manifest.reader_feature_flags &= !feature_flags::FLAG_SEMANTIC_TYPES; + manifest.writer_feature_flags &= !feature_flags::FLAG_SEMANTIC_TYPES; + manifest.version += 1; + write_manifest_file( + dataset.object_store.as_ref(), + dataset.commit_handler.as_ref(), + &dataset.base, + &mut manifest, + None, + &ManifestWriteConfig { + auto_set_feature_flags: false, + ..Default::default() + }, + dataset.manifest_location.naming_scheme, + None, + false, + ) + .await + .unwrap(); + + let err = write_single_column( + uri, + string_array(&DataType::LargeUtf8, &[Some("b")]), + WriteMode::Append, + None, + ) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("should have type string but type was large_string"), + "{err}" + ); + let dataset = write_single_column( + uri, + string_array(&DataType::Utf8, &[Some("b")]), + WriteMode::Append, + None, + ) + .await + .unwrap(); + assert_eq!( + dataset.manifest.reader_feature_flags & feature_flags::FLAG_SEMANTIC_TYPES, + 0 + ); + assert_eq!( + dataset.manifest.writer_feature_flags & feature_flags::FLAG_SEMANTIC_TYPES, + 0 + ); +} + +/// Changing a column's output encoding is a metadata-only schema update. A +/// value that is not an output encoding of the column's type is rejected, and +/// removing the entry restores the type's default. +#[tokio::test] +async fn test_semantic_output_encoding_update() { + use lance_core::datatypes::OUTPUT_ENCODING_META_KEY; + + let dir = TempStrDir::default(); + let mut dataset = write_single_column( + dir.as_str(), + string_array(&DataType::Utf8, &[Some("a")]), + WriteMode::Create, + Some(LanceFileVersion::V2_3), + ) + .await + .unwrap(); + let read_type = |dataset: &Dataset| { + let dataset = dataset.clone(); + async move { + dataset + .scan() + .try_into_batch() + .await + .unwrap() + .column(0) + .data_type() + .clone() + } + }; + + dataset + .update_field_metadata() + .update("s", [(OUTPUT_ENCODING_META_KEY, "large_utf8")]) + .unwrap() + .await + .unwrap(); + assert_eq!(read_type(&dataset).await, DataType::LargeUtf8); + let reopened = Dataset::open(dir.as_str()).await.unwrap(); + assert_eq!(read_type(&reopened).await, DataType::LargeUtf8); + + for invalid in ["large_binary", "utf16"] { + let err = dataset + .update_field_metadata() + .update("s", [(OUTPUT_ENCODING_META_KEY, invalid)]) + .unwrap() + .await + .unwrap_err(); + assert!(err.to_string().contains(invalid), "{err}"); + } + + dataset + .update_field_metadata() + .update("s", [(OUTPUT_ENCODING_META_KEY, None::<&str>)]) + .unwrap() + .await + .unwrap(); + assert_eq!(read_type(&dataset).await, DataType::Utf8); + + // Creating a column sets the entry too, and is rejected the same way. + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("s", DataType::Utf8, true).with_metadata(HashMap::from([( + OUTPUT_ENCODING_META_KEY.to_string(), + "large_binary".to_string(), + )])), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![string_array(&DataType::Utf8, &[None])]).unwrap(); + let other = TempStrDir::default(); + let err = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + other.as_str(), + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }), + ) + .await + .unwrap_err(); + assert!(err.to_string().contains("large_binary"), "{err}"); +} + +/// Merge insert matches a source column by semantic type, so a source in +/// another layout of the column's type updates and inserts rows. +#[tokio::test] +async fn test_semantic_merge_insert_accepts_other_layout() { + use crate::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched}; + + let dir = TempStrDir::default(); + let batch = |ids: Vec, layout: &DataType, values: &[Option<&str>]| { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("s", layout.clone(), true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)), + string_array(layout, values), + ], + ) + .unwrap() + }; + let initial = batch(vec![1, 2], &DataType::Utf8, &[Some("a"), Some("b")]); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(initial.clone())], initial.schema()), + dir.as_str(), + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }), + ) + .await + .unwrap(); + + let source = batch(vec![2, 3], &DataType::LargeUtf8, &[Some("B"), Some("c")]); + let (dataset, _) = MergeInsertBuilder::try_new(Arc::new(dataset), vec!["id".to_string()]) + .unwrap() + .when_matched(WhenMatched::UpdateAll) + .when_not_matched(WhenNotMatched::InsertAll) + .try_build() + .unwrap() + .execute(lance_datafusion::utils::reader_to_stream(Box::new( + RecordBatchIterator::new(vec![Ok(source.clone())], source.schema()), + ))) + .await + .unwrap(); + + let mut scan = dataset.scan(); + scan.order_by(Some(vec![ + crate::dataset::scanner::ColumnOrdering::asc_nulls_first("id".to_string()), + ])) + .unwrap(); + let result = scan.try_into_batch().await.unwrap(); + assert_eq!( + result, + batch( + vec![1, 2, 3], + &DataType::Utf8, + &[Some("a"), Some("B"), Some("c")] + ) + ); +} + +/// Readers of a table under the semantic type contract ask every data file for +/// the table's output layouts, which 2.0 decoders cannot produce, so such a +/// table cannot reference 2.0 files even with mixed file versions enabled. +#[tokio::test] +async fn test_semantic_table_rejects_v2_0_data_files() { + let dir = TempStrDir::default(); + let dataset = write_single_column( + dir.as_str(), + string_array(&DataType::Utf8, &[Some("a")]), + WriteMode::Create, + Some(LanceFileVersion::V2_3), + ) + .await + .unwrap(); + let mut manifest = dataset.manifest.as_ref().clone(); + manifest.reader_feature_flags |= feature_flags::FLAG_MIXED_DATA_FILE_VERSIONS; + manifest.writer_feature_flags |= feature_flags::FLAG_MIXED_DATA_FILE_VERSIONS; + crate::dataset::versions::check_manifest_storage_version(&mut manifest).unwrap(); + + let fragments = Arc::make_mut(&mut manifest.fragments); + fragments[0].files[0].file_major_version = 2; + fragments[0].files[0].file_minor_version = 0; + let err = crate::dataset::versions::check_manifest_storage_version(&mut manifest).unwrap_err(); + assert!(err.to_string().contains("2.1 or later"), "{err}"); +} diff --git a/rust/lance/src/dataset/utils.rs b/rust/lance/src/dataset/utils.rs index 254718464e9..31273aa96c0 100644 --- a/rust/lance/src/dataset/utils.rs +++ b/rust/lance/src/dataset/utils.rs @@ -196,6 +196,60 @@ fn downcast_view_columns( ) } +fn contains_view(data_type: &DataType) -> bool { + match data_type { + DataType::Utf8View | DataType::BinaryView => true, + DataType::List(item) + | DataType::LargeList(item) + | DataType::FixedSizeList(item, _) + | DataType::Map(item, _) => contains_view(item.data_type()), + DataType::Struct(fields) => fields.iter().any(|field| contains_view(field.data_type())), + _ => false, + } +} + +/// Cast columns holding view arrays to the layout `file_schema` records for +/// them. Encoders do not write views; the file records the offset layout a +/// view is stored as. +fn cast_view_columns( + batch: RecordBatch, + file_schema: &ArrowSchema, +) -> std::result::Result { + let schema = batch.schema(); + if !schema + .fields() + .iter() + .any(|field| contains_view(field.data_type())) + { + return Ok(batch); + } + let mut fields = Vec::with_capacity(schema.fields().len()); + let mut columns = Vec::with_capacity(batch.num_columns()); + for (field, column) in schema.fields().iter().zip(batch.columns()) { + let file_type = file_schema + .field_with_name(field.name()) + .ok() + .map(|file_field| file_field.data_type()) + .filter(|file_type| { + contains_view(field.data_type()) && *file_type != field.data_type() + }); + if let Some(file_type) = file_type { + columns.push(arrow_cast::cast(column.as_ref(), file_type)?); + fields.push(field.as_ref().clone().with_data_type(file_type.clone())); + } else { + columns.push(column.clone()); + fields.push(field.as_ref().clone()); + } + } + RecordBatch::try_new( + Arc::new(ArrowSchema::new_with_metadata( + fields, + schema.metadata().clone(), + )), + columns, + ) +} + /// Converts between the Arrow representations callers use and the ones Lance /// stores: Arrow JSON text ↔ Lance JSONB, and top-level view arrays → offset /// arrays. @@ -228,13 +282,21 @@ impl SchemaAdapter { schema.fields().iter().any(|field| has_json_fields(field)) } - pub fn to_physical_batch(&self, batch: RecordBatch) -> Result { - if self.requires_physical_conversion() { - let batch = convert_json_columns(&batch)?; - Ok(downcast_view_columns(&batch)?) + /// Convert a batch to the representation a data file with `file_schema` + /// stores: Arrow JSON text becomes Lance JSONB, and view arrays take the + /// offset layout the file records for them. + pub fn to_file_batch(batch: RecordBatch, file_schema: &ArrowSchema) -> Result { + let batch = if batch + .schema() + .fields() + .iter() + .any(|field| has_arrow_json_fields(field)) + { + convert_json_columns(&batch)? } else { - Ok(batch) - } + batch + }; + Ok(cast_view_columns(batch, file_schema)?) } /// Build the physical Arrow schema for `logical_schema`: Arrow JSON fields diff --git a/rust/lance/src/dataset/versions/mod.rs b/rust/lance/src/dataset/versions/mod.rs index 051e933a6bf..7ec0ec5eb61 100644 --- a/rust/lance/src/dataset/versions/mod.rs +++ b/rust/lance/src/dataset/versions/mod.rs @@ -34,7 +34,7 @@ use lance_file::{ use lance_index::scalar::seed::IndexSeedWriter; use lance_io::object_store::ObjectStore; use lance_io::traits::Writer as ObjectWriter; -use lance_table::feature_flags::FLAG_MIXED_DATA_FILE_VERSIONS; +use lance_table::feature_flags::{FLAG_MIXED_DATA_FILE_VERSIONS, creates_semantic_types}; use lance_table::format::{DataFile, DataStorageFormat, Fragment, Manifest}; use object_store::path::Path; @@ -49,7 +49,7 @@ use super::schema_evolution::optimize::{ ChainedNewColumnTransformOptimizer, SqlToAllNullsOptimizer, }; use super::statistics::FieldStatistics; -use super::write::{self, GenericWriter, TargetBaseInfo, WriteParams, WriterOptions}; +use super::write::{self, GenericWriter, TargetBaseInfo, WriteMode, WriteParams, WriterOptions}; use crate::io::exec::filtered_read::{FilteredReadExec, FilteredReadOptions}; use crate::io::exec::{ AddRowAddrExec, FilterPlan as ExprFilterPlan, LanceScanConfig, LanceStream, TakeExec, @@ -151,6 +151,17 @@ pub async fn write_fragments( } _ => normalized_schema, }; + let semantic_types = dataset.map_or_else( + || creates_semantic_types(version), + |dataset| dataset.manifest.uses_semantic_types(), + ); + if semantic_types && (dataset.is_none() || matches!(params.mode, WriteMode::Overwrite)) { + // This write defines the table schema, so reject inputs that have no + // semantic type or that set an invalid output encoding before any data + // file is written. + normalized_schema.check_output_encoding_entries(None)?; + normalized_schema.to_canonical_types()?; + } let version_name = format!("{version:?}"); let schema = write::prepare_write_schema( dataset, @@ -158,6 +169,15 @@ pub async fn write_fragments( ¶ms, schema_compare_options(version), )?; + // Under the semantic type contract, each data file records the layouts of + // the data written to it, which callers may pass in any accepted layout. + let schema = if semantic_types { + let mut input_schema = Schema::try_from(data.schema().as_ref())?; + input_schema.record_view_layouts(data.schema().as_ref()); + schema.with_input_layouts(&input_schema)? + } else { + schema + }; match version { ConcreteFileVersion::V1 | ConcreteFileVersion::V2_0 | ConcreteFileVersion::V2_1 => { write::validate_legacy_blob_write_schema(&schema, &version_name)?; @@ -354,6 +374,7 @@ fn check_manifest_storage_contract( let mut saw_v1 = false; let mut saw_v2 = false; + let mut first_v2_0 = None; let mut first_file_version = None; let mut first_mismatch = None; let mut first_non_default = None; @@ -363,6 +384,9 @@ fn check_manifest_storage_contract( for fragment in manifest.fragments.iter() { for data_file in fragment.referenced_lance_files() { let file_version = data_file.file_version()?; + if file_version == ConcreteFileVersion::V2_0 && first_v2_0.is_none() { + first_v2_0 = Some(data_file.path.clone()); + } match file_version { ConcreteFileVersion::V1 => saw_v1 = true, ConcreteFileVersion::V2_0 @@ -435,6 +459,15 @@ fn check_manifest_storage_contract( "Dataset snapshot mixes V1 and V2 data files", )); } + // Readers of these tables ask each file for the table's output layouts, + // which 2.0 decoders cannot produce from another stored layout. + if manifest.uses_semantic_types() + && let Some(path) = first_v2_0 + { + return Err(Error::invalid_input(format!( + "Data file '{path}' has version 2.0, but tables that follow the semantic type contract need data file version 2.1 or later" + ))); + } if mixed_enabled && saw_v1 { return Err(Error::invalid_input( diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index d57a238bc2d..32a2fbd1183 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use arrow_array::{ArrayRef, RecordBatch}; +use arrow_schema::{Schema as ArrowSchema, SchemaRef as ArrowSchemaRef}; use bytes::Bytes; use chrono::TimeDelta; use datafusion::physical_plan::SendableRecordBatchStream; @@ -204,15 +205,17 @@ impl Dataset { .parts_dir(&self.data_file_dir_for_base(target.base_id)?) .join(file_name.as_str()); let store = self.object_store(target.base_id).await?; + let file_schema = target.schema.to_data_file_schema()?; let mut writer = V2WriterAdapter::new( file_versions::create_writer( target.version, store.create(&path).await?, - target.schema.as_ref().clone(), + file_schema.clone(), FileWriterOptions::default(), )?, None, preprocessor, + &file_schema, ); let mut data = Box::pin(data); let write_result = async { @@ -940,6 +943,7 @@ where // Keep a copy so failure paths can clean up files written to target bases. let cleanup_bases = target_bases_info.clone(); + let file_schema = ArrowSchema::from(&schema.to_data_file_schema()?); let file_writer_options = params.file_writer_options.clone().unwrap_or_default(); let writer_generator = WriterGenerator::new( object_store.clone(), @@ -1003,7 +1007,7 @@ where // writer then finds nothing left to convert. let batch_chunk = batch_chunk .into_iter() - .map(|batch| SchemaAdapter::new(batch.schema()).to_physical_batch(batch)) + .map(|batch| SchemaAdapter::to_file_batch(batch, &file_schema)) .collect::>>()?; if writer.is_none() { @@ -1802,6 +1806,15 @@ pub(super) fn prepare_write_schema( &normalized_converted_schema, dataset.schema(), )?; + if dataset.manifest.uses_semantic_types() { + // Representation-only inputs are written in their own layout, and + // each data file records it; the table schema does not change. + normalized_converted_schema + .check_compatible(dataset.schema(), &schema_compare_options)?; + return dataset + .schema() + .project_for_write(&normalized_converted_schema); + } if normalized_converted_schema .check_compatible(dataset.schema(), &schema_compare_options) .is_ok() @@ -2031,26 +2044,31 @@ pub(in crate::dataset) struct V2WriterAdapter { writer: current_writer::FileWriter, data_file: Option, preprocessor: Option, + /// The Arrow schema the data file records, which batches are converted to. + file_schema: ArrowSchemaRef, } impl V2WriterAdapter { /// `data_file` describes the file being written and is completed by /// [`GenericWriter::finish`]. Writers that describe their output - /// themselves pass `None` and use [`Self::finish_file`]. + /// themselves pass `None` and use [`Self::finish_file`]. `file_schema` is + /// the schema `writer` records, as [`Schema::to_data_file_schema`] makes it. pub(in crate::dataset) fn new( writer: current_writer::FileWriter, data_file: Option, preprocessor: Option, + file_schema: &Schema, ) -> Self { Self { writer, data_file, preprocessor, + file_schema: Arc::new(ArrowSchema::from(file_schema)), } } pub(in crate::dataset) async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> { - let batch = SchemaAdapter::new(batch.schema()).to_physical_batch(batch.clone())?; + let batch = SchemaAdapter::to_file_batch(batch.clone(), &self.file_schema)?; let batch = match self.preprocessor.as_mut() { Some(pre) => pre.preprocess_batch(&batch).await?, None => batch, @@ -2215,6 +2233,7 @@ where } = options; let (_data_file_key, filename, _data_dir, full_path) = prepare_data_file_path(base_dir, add_data_dir); + let schema = schema.to_data_file_schema()?; let writer = object_store.create(&full_path).await?; let (file_writer, data_file) = create_file_writer( writer, @@ -2227,6 +2246,7 @@ where file_writer, Some(data_file), None, + &schema, ))) } @@ -2259,6 +2279,7 @@ where } = options; let (data_file_key, filename, data_dir, full_path) = prepare_data_file_path(base_dir, add_data_dir); + let schema = &schema.to_data_file_schema()?; let writer = object_store.create(&full_path).await?; let (file_writer, data_file) = create_file_writer( writer, @@ -2283,6 +2304,7 @@ where file_writer, Some(data_file), Some(preprocessor), + schema, ))) } diff --git a/rust/lance/src/dataset/write/insert.rs b/rust/lance/src/dataset/write/insert.rs index 073cd235ee8..a5e385da671 100644 --- a/rust/lance/src/dataset/write/insert.rs +++ b/rust/lance/src/dataset/write/insert.rs @@ -327,6 +327,13 @@ impl<'a> InsertBuilder<'a> { context.storage_version ))); } + if dataset.manifest.uses_semantic_types() + && context.storage_version == ConcreteFileVersion::V2_0 + { + return Err(Error::invalid_input( + "Cannot append data files in version 2.0: tables that follow the semantic type contract need data file version 2.1 or later", + )); + } let mut schema_cmp_opts = crate::dataset::versions::schema_compare_options(version); schema_cmp_opts.compare_nullability = NullabilityComparison::Ignore; schema_cmp_opts.allow_missing_if_nullable = true; @@ -530,11 +537,18 @@ mod test { data_storage_version: Some(target_version), ..Default::default() }; - let dataset = InsertBuilder::new(Arc::new(dataset)) + let result = InsertBuilder::new(Arc::new(dataset)) .with_params(&explicit_params) .execute(vec![batch.clone()]) - .await - .unwrap(); + .await; + // A 2.3 table follows the semantic type contract, whose readers need + // data files of version 2.1 or later. + if default_version == LanceFileVersion::V2_3 && target_version == LanceFileVersion::V2_0 { + let err = result.unwrap_err(); + assert!(err.to_string().contains("2.1 or later"), "{err}"); + return; + } + let dataset = result.unwrap(); assert_eq!( dataset.manifest.data_storage_format.lance_file_format(),