diff --git a/java/lance-jni/src/fragment.rs b/java/lance-jni/src/fragment.rs index 59ce5e553ff..d6255aa2fdb 100644 --- a/java/lance-jni/src/fragment.rs +++ b/java/lance-jni/src/fragment.rs @@ -576,7 +576,7 @@ fn inner_encode_row_ids(env: &mut JNIEnv, row_ids: &JLongArray) -> Result JObject::null(), }; let base_id = convert_to_java_integer(env, self.base_id)?; + let blob_bytes = JLance( + self.blob_bytes + .iter() + .map(|v| *v as i64) + .collect::>(), + ) + .into_java(env)?; Ok(env.new_object( DATA_FILE_CLASS, DATA_FILE_CONSTRUCTOR_SIG, @@ -646,6 +653,7 @@ impl IntoJava for &DataFile { JValueGen::Int(self.file_minor_version as i32), JValueGen::Object(&file_size_bytes), JValueGen::Object(&base_id), + JValueGen::Object(&blob_bytes), ], )?) } @@ -904,6 +912,12 @@ impl FromJObjectWithEnv for JObject<'_> { let file_size_bytes = file_size_bytes.map_or(Default::default(), |r| CachedFileSize::new(r as u64)); let base_id = get_base_id(env, self)?; + let blob_bytes_obj = env.call_method(self, "getBlobBytes", "()[J", &[])?.l()?; + let blob_bytes: Vec = if blob_bytes_obj.is_null() { + Vec::new() + } else { + JLongArray::from(blob_bytes_obj).extract_object(env)? + }; Ok(DataFile { path, fields: fields.into(), @@ -912,6 +926,7 @@ impl FromJObjectWithEnv for JObject<'_> { file_minor_version, file_size_bytes, base_id, + blob_bytes: Arc::from(blob_bytes), }) } } diff --git a/java/lance-jni/src/traits.rs b/java/lance-jni/src/traits.rs index 5a8c9aba6d1..561d3d253bc 100644 --- a/java/lance-jni/src/traits.rs +++ b/java/lance-jni/src/traits.rs @@ -309,6 +309,15 @@ impl FromJObjectWithEnv> for JLongArray<'_> { } } +impl FromJObjectWithEnv> for JLongArray<'_> { + fn extract_object(&self, env: &mut JNIEnv<'_>) -> Result> { + let len = env.get_array_length(self)?; + let mut ret: Vec = vec![0; len as usize]; + env.get_long_array_region(self, 0, ret.as_mut_slice())?; + Ok(ret.into_iter().map(|val| val as u64).collect()) + } +} + impl FromJObjectWithEnv for JObject<'_> { fn extract_object(&self, env: &mut JNIEnv<'_>) -> Result { let ret = env.call_method(self, "intValue", "()I", &[])?.i()?; diff --git a/java/src/main/java/org/lance/fragment/DataFile.java b/java/src/main/java/org/lance/fragment/DataFile.java index 1120a5286c1..d836acc59a7 100644 --- a/java/src/main/java/org/lance/fragment/DataFile.java +++ b/java/src/main/java/org/lance/fragment/DataFile.java @@ -29,6 +29,7 @@ public class DataFile implements Serializable { private final int fileMinorVersion; private final Long fileSizeBytes; private final Integer baseId; + private final long[] blobBytes; public DataFile( String path, @@ -38,6 +39,20 @@ public DataFile( int fileMinorVersion, Long fileSizeBytes, Integer baseId) { + this( + path, fields, columnIndices, fileMajorVersion, fileMinorVersion, fileSizeBytes, baseId, + null); + } + + public DataFile( + String path, + int[] fields, + int[] columnIndices, + int fileMajorVersion, + int fileMinorVersion, + Long fileSizeBytes, + Integer baseId, + long[] blobBytes) { this.path = path; this.fields = fields; this.columnIndices = columnIndices; @@ -45,6 +60,7 @@ public DataFile( this.fileMinorVersion = fileMinorVersion; this.fileSizeBytes = fileSizeBytes; this.baseId = baseId; + this.blobBytes = blobBytes == null ? new long[0] : blobBytes; } public String getPath() { @@ -75,6 +91,16 @@ public Optional getBaseId() { return Optional.ofNullable(baseId); } + /** + * Returns the total size in bytes of the blob payloads backing each field of this file. + * + *

An empty array means blob payload sizes were not recorded (unknown). When non-empty, the + * array has exactly one entry per entry in {@link #getFields()}. + */ + public long[] getBlobBytes() { + return blobBytes; + } + @Override public boolean equals(Object o) { if (this == o) return true; @@ -85,7 +111,8 @@ public boolean equals(Object o) { && Objects.equals(path, that.path) && Arrays.equals(fields, that.fields) && Arrays.equals(columnIndices, that.columnIndices) - && Objects.equals(fileSizeBytes, that.fileSizeBytes); + && Objects.equals(fileSizeBytes, that.fileSizeBytes) + && Arrays.equals(blobBytes, that.blobBytes); } @Override @@ -98,6 +125,7 @@ public String toString() { .add("fileMinorVersion", fileMinorVersion) .add("fileSizeBytes", fileSizeBytes) .add("baseId", baseId) + .add("blobBytes", blobBytes) .toString(); } } diff --git a/protos/table.proto b/protos/table.proto index 9a64230f40f..60f039662d4 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -445,6 +445,21 @@ message DataFile { // The base path index of the data file. Used when the file is imported or referred from another dataset. // Lance use it as key of the base_paths field in Manifest to determine the actual base path of the data file. optional uint32 base_id = 7; + + // Blob payload bytes stored by this data file for each field in `fields`, + // outside the field's regular column pages. This covers blob v2 payloads + // written inline (out-of-line buffers inside this file) and payloads spilled + // to packed or dedicated sidecar `.blob` files. Bytes of external blobs + // (foreign URIs not owned by Lance) are never included. + // + // Recording the sizes here allows reporting of blob storage sizes without + // extra IO, mirroring `IndexFile.size_bytes`. + // + // When present, there must be one entry per entry in `fields` (zero for + // non-blob fields). When empty, blob payload sizes were not recorded — + // either the file was written before this field existed or the writer could + // not know them (e.g. externally created data files). + repeated uint64 blob_bytes = 8; } // DataFile // An overlay file supplies new values for a subset of (row offset, field) cells diff --git a/python/python/lance/fragment.py b/python/python/lance/fragment.py index bca62e2b280..698db3ff637 100644 --- a/python/python/lance/fragment.py +++ b/python/python/lance/fragment.py @@ -225,6 +225,10 @@ class DataFile: The minor version of the data storage format. file_size_bytes : Optional[int] The size of the data file in bytes, if available. + blob_bytes : List[int] + The blob payload bytes stored by this data file, one entry per entry + in `fields` (0 for non-blob fields). Empty if blob payload sizes were + not recorded. """ _path: str @@ -234,6 +238,7 @@ class DataFile: file_minor_version: int = 0 file_size_bytes: Optional[int] = None base_id: Optional[int] = None + blob_bytes: List[int] = field(default_factory=list) def __init__( self, @@ -244,6 +249,7 @@ def __init__( file_minor_version: int = 0, file_size_bytes: Optional[int] = None, base_id: Optional[int] = None, + blob_bytes: List[int] = None, ): # TODO: only we eliminate the path method, we can remove this self._path = path @@ -253,6 +259,7 @@ def __init__( self.file_minor_version = file_minor_version self.file_size_bytes = file_size_bytes self.base_id = base_id + self.blob_bytes = blob_bytes or [] def __repr__(self): # pretend we have a 'path' attribute diff --git a/python/src/fragment.rs b/python/src/fragment.rs index 6b832870280..94dba24404d 100644 --- a/python/src/fragment.rs +++ b/python/src/fragment.rs @@ -894,6 +894,12 @@ impl FromPyObject<'_, '_> for PyLance { let file_size_bytes = CachedFileSize::new(file_size_bytes.unwrap_or(0)); let fields: Vec = ob.getattr("fields")?.extract()?; let column_indices: Vec = ob.getattr("column_indices")?.extract()?; + // Older DataFile objects may not have a blob_bytes attribute; treat + // that the same as "blob payload sizes not recorded" (empty). + let blob_bytes: Vec = ob + .getattr("blob_bytes") + .and_then(|v| v.extract()) + .unwrap_or_default(); Ok(Self(DataFile { path: ob.getattr("path")?.extract()?, fields: fields.into(), @@ -902,6 +908,7 @@ impl FromPyObject<'_, '_> for PyLance { file_minor_version: ob.getattr("file_minor_version")?.extract()?, file_size_bytes, base_id: ob.getattr("base_id")?.extract()?, + blob_bytes: Arc::from(blob_bytes), })) } } @@ -926,6 +933,7 @@ impl<'py> IntoPyObject<'py> for PyLance<&DataFile> { self.0.file_minor_version, file_size_bytes, self.0.base_id, + self.0.blob_bytes.to_vec(), )) } } diff --git a/rust/lance-table/benches/manifest_intern.rs b/rust/lance-table/benches/manifest_intern.rs index 81bd57c1a22..d943f828c0e 100644 --- a/rust/lance-table/benches/manifest_intern.rs +++ b/rust/lance-table/benches/manifest_intern.rs @@ -58,6 +58,7 @@ fn make_uniform_pb_fragments(n: u64, num_fields: usize) -> Vec file_minor_version: 0, file_size_bytes: 0, base_id: None, + blob_bytes: vec![], }], overlays: vec![], deletion_file: None, @@ -135,6 +136,7 @@ fn make_diverse_pb_fragments( file_minor_version: 0, file_size_bytes: 0, base_id: None, + blob_bytes: vec![], }], overlays: vec![], deletion_file: None, diff --git a/rust/lance-table/src/format/fragment.rs b/rust/lance-table/src/format/fragment.rs index 3a6f4269b0e..9cce4281281 100644 --- a/rust/lance-table/src/format/fragment.rs +++ b/rust/lance-table/src/format/fragment.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use std::num::NonZero; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use lance_core::Error; use lance_core::deepsize::DeepSizeOf; @@ -54,13 +54,21 @@ pub struct DataFile { /// The base path of the datafile, when the datafile is outside the dataset. pub base_id: Option, + + /// Blob v2 payload bytes stored for each field in `fields`, outside the + /// field's regular column pages (inline out-of-line buffers in this file + /// plus packed/dedicated sidecar `.blob` files). External blob bytes are + /// never included. Empty when not recorded: files written before this + /// field existed, or data files whose sidecar sizes the writer could not + /// know. When non-empty, has exactly one entry per entry in `fields`. + pub blob_bytes: Arc<[u64]>, } // Custom Serialize: convert Arc<[i32]> to slice for transparent JSON output impl Serialize for DataFile { fn serialize(&self, serializer: S) -> std::result::Result { use serde::ser::SerializeStruct; - let mut s = serializer.serialize_struct("DataFile", 7)?; + let mut s = serializer.serialize_struct("DataFile", 8)?; s.serialize_field("path", &self.path)?; s.serialize_field("fields", self.fields.as_ref())?; s.serialize_field("column_indices", self.column_indices.as_ref())?; @@ -68,6 +76,7 @@ impl Serialize for DataFile { s.serialize_field("file_minor_version", &self.file_minor_version)?; s.serialize_field("file_size_bytes", &self.file_size_bytes)?; s.serialize_field("base_id", &self.base_id)?; + s.serialize_field("blob_bytes", self.blob_bytes.as_ref())?; s.end() } } @@ -87,6 +96,8 @@ impl<'de> Deserialize<'de> for DataFile { file_minor_version: u32, file_size_bytes: CachedFileSize, base_id: Option, + #[serde(default)] + blob_bytes: Vec, } let helper = DataFileHelper::deserialize(deserializer)?; @@ -98,10 +109,25 @@ impl<'de> Deserialize<'de> for DataFile { file_minor_version: helper.file_minor_version, file_size_bytes: helper.file_size_bytes, base_id: helper.base_id, + blob_bytes: blob_bytes_arc(helper.blob_bytes), }) } } +/// Shared allocation for the common "blob payload sizes not recorded" case, so +/// the many empty tallies across a manifest don't each allocate. +static EMPTY_BLOB_BYTES: LazyLock> = LazyLock::new(|| Arc::from([])); + +/// Convert a blob byte tally to `Arc<[u64]>`, collapsing the common empty case +/// to a single shared allocation. +fn blob_bytes_arc(blob_bytes: Vec) -> Arc<[u64]> { + if blob_bytes.is_empty() { + EMPTY_BLOB_BYTES.clone() + } else { + Arc::from(blob_bytes) + } +} + impl DataFile { /// Create a `DataFile` and encode its exact format version for manifest storage. pub fn new( @@ -121,9 +147,18 @@ impl DataFile { file_minor_version, file_size_bytes: file_size_bytes.into(), base_id, + blob_bytes: EMPTY_BLOB_BYTES.clone(), } } + /// Record blob v2 payload bytes per field in `fields`. Pass an empty vec + /// (or don't call) when sizes are unknown; when non-empty, `blob_bytes` + /// must have one entry per entry in `fields`. + pub fn with_blob_bytes(mut self, blob_bytes: Vec) -> Self { + self.blob_bytes = blob_bytes_arc(blob_bytes); + self + } + /// Create a new `DataFile` whose fields and column indices will be set later. pub fn new_unstarted(path: impl Into, file_version: ConcreteFileVersion) -> Self { let (file_major_version, file_minor_version) = file_version.to_data_file_numbers(); @@ -135,6 +170,7 @@ impl DataFile { file_minor_version, file_size_bytes: Default::default(), base_id: None, + blob_bytes: EMPTY_BLOB_BYTES.clone(), } } @@ -196,6 +232,16 @@ impl DataFile { "contained fewer column_indices than fields", )); } + if !self.blob_bytes.is_empty() && self.blob_bytes.len() != self.fields.len() { + return Err(Error::corrupt_file( + base_path.clone().join(self.path.clone()), + format!( + "contained {} blob_bytes entries for {} fields", + self.blob_bytes.len(), + self.fields.len() + ), + )); + } Ok(()) } } @@ -210,6 +256,7 @@ impl From<&DataFile> for pb::DataFile { file_minor_version: df.file_minor_version, file_size_bytes: df.file_size_bytes.get().map_or(0, |v| v.get()), base_id: df.base_id, + blob_bytes: df.blob_bytes.to_vec(), } } } @@ -226,6 +273,7 @@ impl TryFrom for DataFile { file_minor_version: proto.file_minor_version, file_size_bytes: CachedFileSize::new(proto.file_size_bytes), base_id: proto.base_id, + blob_bytes: blob_bytes_arc(proto.blob_bytes), }) } } @@ -347,6 +395,10 @@ impl DataFileFieldInterner { file_minor_version: proto.file_minor_version, file_size_bytes: CachedFileSize::new(proto.file_size_bytes), base_id: proto.base_id, + // Non-empty blob byte tallies rarely repeat across files, so they + // are not worth interning; the common empty case already collapses + // to a single shared allocation. + blob_bytes: blob_bytes_arc(proto.blob_bytes), }) } @@ -968,9 +1020,9 @@ mod tests { json!({ "id": 123, "files":[ - {"path": "foobar.lance", "fields": [0], "column_indices": [], + {"path": "foobar.lance", "fields": [0], "column_indices": [], "file_major_version": MAJOR_VERSION, "file_minor_version": MINOR_VERSION, - "file_size_bytes": null, "base_id": null } + "file_size_bytes": null, "base_id": null, "blob_bytes": [] } ], "deletion_file": {"read_version": 123, "id": 456, "file_type": "array", "num_deleted_rows": 10, "base_id": null}, @@ -981,6 +1033,41 @@ mod tests { assert_eq!(fragment, frag2); } + #[test] + fn data_file_blob_bytes_roundtrip_and_validate() { + let data_file = DataFile::new( + "foo.lance", + vec![1, 2, 3], + vec![0, -1, 1], + ConcreteFileVersion::V2_2, + None, + None, + ) + .with_blob_bytes(vec![0, 42, 0]); + + let proto = pb::DataFile::from(&data_file); + assert_eq!(proto.blob_bytes, vec![0, 42, 0]); + let from_proto = DataFile::try_from(proto).unwrap(); + assert_eq!(from_proto, data_file); + + let base_path = Path::from("base"); + data_file.validate(&base_path).unwrap(); + + // Old writers never recorded blob bytes; an empty tally must validate. + let unknown = from_proto.with_blob_bytes(vec![]); + unknown.validate(&base_path).unwrap(); + + // A recorded tally must line up with `fields` entry for entry. + let misaligned = data_file.with_blob_bytes(vec![42]); + let err = misaligned.validate(&base_path).unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err}"); + assert!( + err.to_string() + .contains("1 blob_bytes entries for 3 fields"), + "{err}" + ); + } + #[test] fn data_file_validate_allows_extra_columns() { let data_file = DataFile { @@ -992,6 +1079,7 @@ mod tests { file_minor_version: MINOR_VERSION as u32, file_size_bytes: Default::default(), base_id: None, + blob_bytes: Arc::from([]), }; let base_path = Path::from("base"); diff --git a/rust/lance/src/blob.rs b/rust/lance/src/blob.rs index b48a5603be6..ca74838e8ad 100644 --- a/rust/lance/src/blob.rs +++ b/rust/lance/src/blob.rs @@ -330,7 +330,22 @@ fn validate_range(offset: u64, size: u64, object_size: u64, label: &str) -> Resu Ok(()) } -fn validate_prepared_blob_value_array(field: &Field, array: &ArrayRef) -> Result<()> { +/// Payload bytes observed while validating a prepared blob v2 array. +#[derive(Debug, Default, Clone, Copy)] +pub(crate) struct PreparedBlobTally { + /// Total bytes of inline payloads. The encoder stores these out-of-line in + /// the data file being written, so they count as bytes stored by that file. + pub inline_bytes: u64, + /// Whether any packed or dedicated rows were present. Their payloads live + /// in sidecar files this writer did not produce, so their sizes cannot be + /// attributed to the file being written. + pub has_sidecar_references: bool, +} + +fn validate_prepared_blob_value_array( + field: &Field, + array: &ArrayRef, +) -> Result { if !is_prepared_blob_v2_field(field) { return Err(blob_v2_shape_error(field)); } @@ -364,6 +379,7 @@ fn validate_prepared_blob_value_array(field: &Field, array: &ArrayRef) -> Result .ok_or_else(|| Error::invalid_input("Prepared blob struct missing `position` field"))? .as_primitive::(); + let mut tally = PreparedBlobTally::default(); for row in 0..struct_arr.len() { if struct_arr.is_null(row) { continue; @@ -381,6 +397,7 @@ fn validate_prepared_blob_value_array(field: &Field, array: &ArrayRef) -> Result "Prepared inline blob row {row} must set `data`" ))); } + tally.inline_bytes += data_col.value_length(row) as u64; } BlobKind::Packed => { if blob_id_col.is_null(row) @@ -399,6 +416,7 @@ fn validate_prepared_blob_value_array(field: &Field, array: &ArrayRef) -> Result "Prepared packed blob row {row} range overflows u64: offset={offset}, size={size}" )) })?; + tally.has_sidecar_references = true; } BlobKind::Dedicated => { if blob_id_col.is_null(row) || blob_size_col.is_null(row) { @@ -407,6 +425,7 @@ fn validate_prepared_blob_value_array(field: &Field, array: &ArrayRef) -> Result ))); } validate_blob_id(blob_id_col.value(row))?; + tally.has_sidecar_references = true; } BlobKind::External => { if uri_col.is_null(row) || uri_col.value(row).is_empty() { @@ -433,11 +452,17 @@ fn validate_prepared_blob_value_array(field: &Field, array: &ArrayRef) -> Result } } - Ok(()) + Ok(tally) } /// Validate a writer-side prepared blob v2 array before it reaches the encoder. -pub(crate) fn validate_prepared_blob_array(field: &Field, array: &ArrayRef) -> Result<()> { +/// +/// Returns the payload byte tally observed during validation so writers can +/// record blob storage sizes without a second pass over the array. +pub(crate) fn validate_prepared_blob_array( + field: &Field, + array: &ArrayRef, +) -> Result { validate_prepared_blob_value_array(field, array) } diff --git a/rust/lance/src/dataset/blob.rs b/rust/lance/src/dataset/blob.rs index 7856fe9cdaf..43046525672 100644 --- a/rust/lance/src/dataset/blob.rs +++ b/rust/lance/src/dataset/blob.rs @@ -285,6 +285,19 @@ pub struct BlobPreprocessor { /// this write job only; it is not persisted into the dataset schema. pack_file_size_override: Option, field_processors: Vec, + /// Blob v2 payload bytes stored so far by the data file being written, + /// keyed by blob field id: payloads this preprocessor routed to inline, + /// packed, or dedicated storage, plus inline payloads of user-supplied + /// prepared descriptors (which the encoder stores out-of-line in the same + /// file). Recorded in the resulting [`DataFile`](lance_table::format::DataFile) + /// so statistics can report blob storage sizes without extra IO. + blob_bytes: HashMap, + /// False when the file stores blob payload bytes this preprocessor cannot + /// attribute: blob fields nested under wrappers it does not traverse + /// (e.g. FixedSizeList or Map), or user-supplied packed/dedicated + /// descriptors whose sidecars it did not write. When false the tally must + /// not be recorded — an unknown tally (empty) beats a wrong one. + is_tally_complete: bool, external_base_resolver: Option>, allow_external_blob_outside_bases: bool, external_blob_mode: ExternalBlobMode, @@ -314,13 +327,23 @@ struct BlobPreprocessField { kind: BlobPreprocessFieldKind, } +/// Per-field settings for ingesting a logical blob v2 column. +#[derive(Clone, Debug)] +struct BlobV2FieldConfig { + field_id: i32, + inline_threshold: usize, + dedicated_threshold: usize, + pack_file_threshold: usize, + writer_metadata: HashMap, +} + #[derive(Clone, Debug)] enum BlobPreprocessFieldKind { - BlobV2 { - inline_threshold: usize, - dedicated_threshold: usize, - pack_file_threshold: usize, - writer_metadata: HashMap, + BlobV2(BlobV2FieldConfig), + /// Already-prepared descriptors: validated and tallied, then passed + /// through unchanged. + PreparedBlobV2 { + field_id: i32, }, Struct { children: Vec, @@ -332,11 +355,24 @@ enum BlobPreprocessFieldKind { } impl BlobPreprocessField { - fn new(field: &ArrowField) -> Result { + /// Classify `field` (paired with its Lance schema counterpart, which + /// carries the field id used for blob byte tallies). + /// + /// Sets `has_unreachable_blob` when a blob v2 field is nested under a + /// wrapper this preprocessor does not traverse (e.g. FixedSizeList or + /// Map): the file may store blob payload bytes the tally cannot see, so + /// no tally must be recorded for it. + fn new( + field: &ArrowField, + lance_field: &LanceField, + has_unreachable_blob: &mut bool, + ) -> Result { if field.is_blob_v2() { if is_prepared_blob_v2_field(field) { return Ok(Self { - kind: BlobPreprocessFieldKind::Passthrough, + kind: BlobPreprocessFieldKind::PreparedBlobV2 { + field_id: lance_field.id, + }, }); } if !is_logical_blob_v2_field(field) { @@ -349,7 +385,8 @@ impl BlobPreprocessField { ))); } return Ok(Self { - kind: BlobPreprocessFieldKind::BlobV2 { + kind: BlobPreprocessFieldKind::BlobV2(BlobV2FieldConfig { + field_id: lance_field.id, inline_threshold: blob_inline_threshold_from_metadata( field.metadata(), field.name(), @@ -363,24 +400,34 @@ impl BlobPreprocessField { field.name(), )?, writer_metadata: field.metadata().clone(), - }, + }), }); } - if let ArrowDataType::Struct(children) = field.data_type() { + if let ArrowDataType::Struct(children) = field.data_type() + && children.len() == lance_field.children.len() + { let children = children .iter() - .map(|child| Self::new(child.as_ref())) + .zip(lance_field.children.iter()) + .map(|(child, lance_child)| { + Self::new(child.as_ref(), lance_child, has_unreachable_blob) + }) .collect::>>()?; if children.iter().any(|child| child.requires_preprocessing()) { return Ok(Self { kind: BlobPreprocessFieldKind::Struct { children }, }); } + return Ok(Self { + kind: BlobPreprocessFieldKind::Passthrough, + }); } - if let ArrowDataType::List(child) | ArrowDataType::LargeList(child) = field.data_type() { - let child = Self::new(child.as_ref())?; + if let ArrowDataType::List(child) | ArrowDataType::LargeList(child) = field.data_type() + && let Some(lance_child) = lance_field.children.first() + { + let child = Self::new(child.as_ref(), lance_child, has_unreachable_blob)?; if child.requires_preprocessing() { return Ok(Self { kind: BlobPreprocessFieldKind::List { @@ -388,8 +435,16 @@ impl BlobPreprocessField { }, }); } + return Ok(Self { + kind: BlobPreprocessFieldKind::Passthrough, + }); } + // Any other wrapper (FixedSizeList, Map, ...) is not traversed; if a + // blob v2 field hides below it, its payload bytes cannot be tallied. + if field_contains_blob_v2(lance_field) { + *has_unreachable_blob = true; + } Ok(Self { kind: BlobPreprocessFieldKind::Passthrough, }) @@ -398,6 +453,21 @@ impl BlobPreprocessField { fn requires_preprocessing(&self) -> bool { !matches!(self.kind, BlobPreprocessFieldKind::Passthrough) } + + /// Collect the ids of blob v2 fields this processor will tally. + fn collect_blob_field_ids(&self, ids: &mut Vec) { + match &self.kind { + BlobPreprocessFieldKind::BlobV2(BlobV2FieldConfig { field_id, .. }) + | BlobPreprocessFieldKind::PreparedBlobV2 { field_id } => ids.push(*field_id), + BlobPreprocessFieldKind::Struct { children } => { + for child in children { + child.collect_blob_field_ids(ids); + } + } + BlobPreprocessFieldKind::List { child } => child.collect_blob_field_ids(ids), + BlobPreprocessFieldKind::Passthrough => {} + } + } } impl ExternalBlobSource { @@ -465,11 +535,30 @@ impl BlobPreprocessor { ) -> Result { let pack_writer = RollingPackedBlobWriter::new(); let arrow_schema = arrow_schema::Schema::from(schema); + if arrow_schema.fields().len() != schema.fields.len() { + return Err(Error::internal(format!( + "Arrow schema has {} fields but Lance schema has {}", + arrow_schema.fields().len(), + schema.fields.len() + ))); + } + let mut has_unreachable_blob = false; let field_processors = arrow_schema .fields() .iter() - .map(|field| BlobPreprocessField::new(field.as_ref())) + .zip(schema.fields.iter()) + .map(|(field, lance_field)| { + BlobPreprocessField::new(field.as_ref(), lance_field, &mut has_unreachable_blob) + }) .collect::>>()?; + // Start every reachable blob field at zero so a file whose blobs are + // all null or external still records a known-zero tally rather than + // "not recorded". + let mut blob_field_ids = Vec::new(); + for processor in &field_processors { + processor.collect_blob_field_ids(&mut blob_field_ids); + } + let blob_bytes = blob_field_ids.into_iter().map(|id| (id, 0)).collect(); Ok(Self { object_store, data_dir, @@ -478,6 +567,8 @@ impl BlobPreprocessor { pack_writer, pack_file_size_override, field_processors, + blob_bytes, + is_tally_complete: !has_unreachable_blob, external_base_resolver, allow_external_blob_outside_bases, external_blob_mode, @@ -648,28 +739,46 @@ impl BlobPreprocessor { field: &'a Arc, ) -> BoxFuture<'a, Result<(ArrayRef, Arc)>> { async move { + // The batch shape wins over the schema classification: on append + // the writer schema is the dataset's logical blob schema, yet a + // batch may still carry already-prepared descriptors that must be + // passed through untouched, not re-ingested as logical blobs. if is_prepared_blob_v2_field(field.as_ref()) { - validate_prepared_blob_array(field.as_ref(), &array)?; + let tally = validate_prepared_blob_array(field.as_ref(), &array)?; + match &processor.kind { + BlobPreprocessFieldKind::BlobV2(BlobV2FieldConfig { field_id, .. }) + | BlobPreprocessFieldKind::PreparedBlobV2 { field_id } => { + // Inline payloads of prepared descriptors are stored + // out-of-line in this file by the encoder; packed and + // dedicated descriptors reference sidecars this + // writer did not produce, so the tally cannot vouch + // for them. + *self.blob_bytes.entry(*field_id).or_insert(0) += tally.inline_bytes; + if tally.has_sidecar_references { + self.is_tally_complete = false; + } + } + _ => { + // No blob field id to attribute these payloads to. + self.is_tally_complete = false; + } + } return Ok((array, field.clone())); } match &processor.kind { BlobPreprocessFieldKind::Passthrough => Ok((array, field.clone())), - BlobPreprocessFieldKind::BlobV2 { - inline_threshold, - dedicated_threshold, - pack_file_threshold, - writer_metadata, - } => { - self.preprocess_blob_array( - array, - field.as_ref(), - *inline_threshold, - *dedicated_threshold, - *pack_file_threshold, - writer_metadata, - ) - .await + BlobPreprocessFieldKind::PreparedBlobV2 { .. } => { + // Classified prepared but the batch is not in prepared + // shape: pass through unchanged (matching pre-tally + // behavior) and leave the tally unrecorded rather than + // guessing. + self.is_tally_complete = false; + Ok((array, field.clone())) + } + BlobPreprocessFieldKind::BlobV2(config) => { + self.preprocess_blob_array(array, field.as_ref(), config) + .await } BlobPreprocessFieldKind::Struct { children } => { self.preprocess_struct_array(array, field.as_ref(), children) @@ -821,11 +930,15 @@ impl BlobPreprocessor { &mut self, array: ArrayRef, field: &ArrowField, - inline_threshold: usize, - dedicated_threshold: usize, - pack_file_threshold: usize, - writer_metadata: &HashMap, + config: &BlobV2FieldConfig, ) -> Result<(ArrayRef, Arc)> { + let &BlobV2FieldConfig { + field_id, + inline_threshold, + dedicated_threshold, + pack_file_threshold, + ref writer_metadata, + } = config; let struct_arr = array .as_any() .downcast_ref::() @@ -847,6 +960,9 @@ impl BlobPreprocessor { .map(|col| col.as_primitive::()); let mut blob_writer = self.blob_writer_with_metadata(field, writer_metadata.clone()); + // Payload bytes routed to storage owned by the file being written + // (inline, packed, or dedicated); external references are excluded. + let mut stored_bytes: u64 = 0; for i in 0..struct_arr.len() { if struct_arr.is_null(i) { @@ -876,6 +992,7 @@ impl BlobPreprocessor { ) .await?; blob_writer.push(value)?; + stored_bytes += data_len as u64; continue; } @@ -887,6 +1004,7 @@ impl BlobPreprocessor { ) .await?; blob_writer.push(value)?; + stored_bytes += data_len as u64; continue; } @@ -921,6 +1039,7 @@ impl BlobPreprocessor { ) .await?; blob_writer.push(value)?; + stored_bytes += data_len; continue; } @@ -929,10 +1048,12 @@ impl BlobPreprocessor { .write_packed(pack_file_threshold, BlobWriteSource::External(&source)) .await?; blob_writer.push(value)?; + stored_bytes += data_len; continue; } let data = source.read_all().await?; + stored_bytes += data.len() as u64; blob_writer.push_inline(data)?; continue; } @@ -961,11 +1082,13 @@ impl BlobPreprocessor { if has_data { blob_writer.push_inline(Bytes::copy_from_slice(data_col.value(i)))?; + stored_bytes += data_len as u64; } else { blob_writer.push_null()?; } } + *self.blob_bytes.entry(field_id).or_insert(0) += stored_bytes; let column = blob_writer.finish()?; let (field, array) = column.into_parts(); Ok((array, Arc::new(field))) @@ -974,6 +1097,18 @@ impl BlobPreprocessor { pub(crate) async fn finish(&mut self) -> Result<()> { self.pack_writer.finish().await } + + /// Blob v2 payload bytes stored by the data file being written, keyed by + /// blob field id (zero entries for blob fields whose payloads were all + /// null or external). + /// + /// Returns `None` when the file may store blob payload bytes this + /// preprocessor could not attribute (see `is_tally_complete`); no tally + /// should be recorded in that case. Empty when the schema has no + /// (reachable) blob v2 fields. + pub(super) fn blob_bytes(&self) -> Option<&HashMap> { + self.is_tally_complete.then_some(&self.blob_bytes) + } } pub async fn preprocess_blob_batches( @@ -987,6 +1122,11 @@ pub async fn preprocess_blob_batches( Ok(out) } +/// Returns true when `field` is, or contains, a blob v2 field. +fn field_contains_blob_v2(field: &LanceField) -> bool { + field.is_blob_v2() || field.children.iter().any(field_contains_blob_v2) +} + /// Mutable state for a [`BlobFile`] cursor. /// /// The cursor is logical to the blob slice, not the backing object. Once closed, @@ -5143,6 +5283,193 @@ mod tests { blobs[1].as_ref().unwrap().read().await.unwrap().as_ref(), b"append" ); + + // Both write paths must record the inline payload bytes in the + // manifest: the logical write stored b"initial" (7 bytes) and the + // prepared append stored b"append" (6 bytes) out-of-line in their + // respective data files. + let blob_field_id = dataset.schema().field("blob").unwrap().id; + let fragments = dataset.fragments(); + for (fragment, expected) in fragments.iter().zip([7u64, 6u64]) { + assert_eq!(recorded_blob_bytes(fragment, blob_field_id), Some(expected)); + } + } + + /// The blob payload bytes recorded for `field_id` in the fragment's sole + /// data file, or `None` when no tally was recorded. + fn recorded_blob_bytes(fragment: &lance_table::format::Fragment, field_id: i32) -> Option { + assert_eq!(fragment.files.len(), 1); + let file = &fragment.files[0]; + if file.blob_bytes.is_empty() { + return None; + } + assert_eq!(file.blob_bytes.len(), file.fields.len()); + Some( + file.fields + .iter() + .zip(file.blob_bytes.iter()) + .filter(|(id, _)| **id == field_id) + .map(|(_, bytes)| *bytes) + .sum(), + ) + } + + #[tokio::test] + async fn test_prepared_sidecar_descriptors_leave_blob_bytes_unrecorded() { + // Packed/dedicated prepared descriptors reference sidecar files this + // write does not produce, so the writer cannot vouch for their sizes: + // the tally must be left unrecorded (empty), not recorded as wrong. + let test_dir = TempStrDir::default(); + let logical_schema = Arc::new(Schema::new(vec![blob_field("blob", true)])); + let mut initial_builder = BlobArrayBuilder::new(1); + initial_builder.push_bytes(b"initial").unwrap(); + let initial_batch = RecordBatch::try_new( + logical_schema.clone(), + vec![initial_builder.finish().unwrap()], + ) + .unwrap(); + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(initial_batch)], logical_schema.clone()), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let sidecar_owner_path = dataset.data_dir().join("pre-written.lance"); + let mut blob_writer = BlobDescriptorArrayBuilder::new("blob"); + let mut packed = + PackedBlobWriter::try_new(dataset.object_store.as_ref().clone(), sidecar_owner_path, 1) + .await + .unwrap(); + packed.write_blob(b"packed-elsewhere").await.unwrap(); + blob_writer.extend(packed.finish().await.unwrap()).unwrap(); + let (prepared_field, prepared_array) = blob_writer.finish().unwrap().into_parts(); + let append_schema = Arc::new(Schema::new(vec![prepared_field])); + let append_batch = + RecordBatch::try_new(append_schema.clone(), vec![prepared_array]).unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(append_batch)], append_schema), + dataset, + Some(WriteParams { + mode: WriteMode::Append, + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + let blob_field_id = dataset.schema().field("blob").unwrap().id; + let fragments = dataset.fragments(); + assert_eq!(recorded_blob_bytes(&fragments[0], blob_field_id), Some(7)); + assert_eq!(recorded_blob_bytes(&fragments[1], blob_field_id), None); + } + + #[tokio::test] + async fn test_blob_field_under_untraversed_wrapper_leaves_blob_bytes_unrecorded() { + // A blob v2 field nested under a wrapper the preprocessor does not + // traverse (FixedSizeList here) stores payload bytes the tally cannot + // see. The whole file's tally must be left unrecorded — recording a + // definite 0 for the wrapped field (while the sibling blob records + // bytes) would silently understate storage forever. + let test_dir = TempStrDir::default(); + let mut wrapped_builder = BlobDescriptorArrayBuilder::new("item"); + wrapped_builder + .push_inline(Bytes::from_static(b"wrapped!")) + .unwrap(); + let (wrapped_field, wrapped_array) = wrapped_builder.finish().unwrap().into_parts(); + let fsl = arrow_array::FixedSizeListArray::try_new( + Arc::new(wrapped_field), + 1, + wrapped_array, + None, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("wrapped", fsl.data_type().clone(), false), + blob_field("blob", true), + ])); + let mut blobs = BlobArrayBuilder::new(1); + blobs.push_bytes(b"top-level-payload").unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(fsl) as ArrayRef, blobs.finish().unwrap()], + ) + .unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + let blob_field_id = dataset.schema().field("blob").unwrap().id; + let fragments = dataset.fragments(); + assert_eq!(recorded_blob_bytes(&fragments[0], blob_field_id), None); + } + + #[tokio::test] + async fn test_sliced_list_of_prepared_blobs_tallies_only_own_rows() { + // The write pipeline zero-copy slices batches to honor + // max_rows_per_file; each slice of a list column still exposes the + // full values buffer. The tally must count each file's own rows only, + // not the whole buffer once per file. + let test_dir = TempStrDir::default(); + let mut item_builder = BlobDescriptorArrayBuilder::new("item"); + item_builder + .push_inline(Bytes::from(vec![1u8; 10])) + .unwrap(); + item_builder + .push_inline(Bytes::from(vec![2u8; 20])) + .unwrap(); + let (item_field, item_array) = item_builder.finish().unwrap().into_parts(); + let list = arrow_array::ListArray::try_new( + Arc::new(item_field), + arrow_buffer::OffsetBuffer::from_lengths([1, 1]), + item_array, + None, + ) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "blobs", + list.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(list) as ArrayRef]).unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + max_rows_per_file: 1, + ..Default::default() + }), + ) + .await + .unwrap(); + + let blob_field_id = dataset + .schema() + .field("blobs") + .unwrap() + .children + .first() + .unwrap() + .id; + let fragments = dataset.fragments(); + assert_eq!(fragments.len(), 2); + assert_eq!(recorded_blob_bytes(&fragments[0], blob_field_id), Some(10)); + assert_eq!(recorded_blob_bytes(&fragments[1], blob_field_id), Some(20)); } #[tokio::test] @@ -5220,6 +5547,30 @@ mod tests { .create_data_file(&data_file_name, None) .await .unwrap(); + // The replacement file owns the packed sidecar holding + // b"nested-replacement"; record its tally and check it survives the + // commit (the old file's tally must not linger on the new path). + let blob_field_id = dataset + .schema() + .field("info") + .unwrap() + .children + .iter() + .find(|child| child.name == "blob") + .unwrap() + .id; + let replacement_tally = data_file + .fields + .iter() + .map(|id| { + if *id == blob_field_id { + b"nested-replacement".len() as u64 + } else { + 0 + } + }) + .collect::>(); + let data_file = data_file.with_blob_bytes(replacement_tally.clone()); let transaction = Transaction { read_version: dataset.manifest.version, uuid: Uuid::new_v4().hyphenated().to_string(), @@ -5236,6 +5587,16 @@ mod tests { .unwrap(), ); + let replaced_file = dataset.fragments()[0] + .files + .iter() + .find(|file| file.path == data_file_name) + .expect("replacement data file should be in the fragment"); + assert_eq!( + replaced_file.blob_bytes.as_ref(), + replacement_tally.as_slice() + ); + let blobs = dataset .take_blobs_by_indices(&[0], "info.blob") .await diff --git a/rust/lance/src/dataset/files.rs b/rust/lance/src/dataset/files.rs index 2214822ed90..35750bb2c8e 100644 --- a/rust/lance/src/dataset/files.rs +++ b/rust/lance/src/dataset/files.rs @@ -1026,6 +1026,7 @@ mod tests { file_minor_version: 0, file_size_bytes: CachedFileSize::unknown(), base_id, + blob_bytes: Arc::from([]), }; let fragment = Fragment { diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 7fcc7423cb2..f5a6a132110 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -6088,6 +6088,7 @@ mod tests { file_minor_version: 1, file_size_bytes: CachedFileSize::unknown(), base_id: None, + blob_bytes: Arc::from([]), }; let full_struct = diff --git a/rust/lance/src/dataset/statistics.rs b/rust/lance/src/dataset/statistics.rs index 627ccfc1081..66d52c4b0a1 100644 --- a/rust/lance/src/dataset/statistics.rs +++ b/rust/lance/src/dataset/statistics.rs @@ -23,6 +23,13 @@ pub struct FieldStatistics { /// Amount of data in the field (after compression, if any) /// /// This will be 0 if the data storage version is less than 2 + /// + /// For blob v2 fields this includes the payload bytes Lance stores outside + /// the field's column pages (inline out-of-line buffers plus packed and + /// dedicated sidecar `.blob` files), as recorded in the manifest at write + /// time. Bytes of external blobs (foreign URIs) are never included, and + /// data files written before the tally existed contribute only their + /// descriptor bytes. pub bytes_on_disk: u64, } @@ -53,6 +60,21 @@ impl DatasetStatisticsExt for Dataset { ) })); if !self.is_legacy_storage() { + // Blob v2 payloads live outside column pages (inline out-of-line + // buffers and sidecar `.blob` files), so file metadata never sees + // them. Writers record their per-field sizes in the manifest + // (`DataFile::blob_bytes`); fold those in here. Files without a + // recorded tally (older writers, externally created data files) + // contribute only their descriptor bytes. + for fragment in self.fragments().iter() { + for file in &fragment.files { + for (field_id, blob_bytes) in file.fields.iter().zip(file.blob_bytes.iter()) { + if let Some(stats) = field_stats.get_mut(&(*field_id as u32)) { + stats.bytes_on_disk += blob_bytes; + } + } + } + } let scan_scheduler = ScanScheduler::new( self.object_store.clone(), SchedulerConfig::max_bandwidth(self.object_store.as_ref()), @@ -182,6 +204,7 @@ impl<'a> DatasetStatistics<'a> { #[cfg(test)] mod tests { + use std::num::NonZeroUsize; use std::sync::Arc; use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchIterator}; @@ -189,6 +212,7 @@ mod tests { use lance_core::utils::tempfile::TempStrDir; use lance_file::version::LanceFileVersion; + use crate::blob::{BlobArrayBuilder, BlobFieldOptions, blob_field_with_options}; use crate::dataset::WriteParams; use super::*; @@ -244,4 +268,155 @@ mod tests { "bytes_on_disk should include the remaining column after drop_columns" ); } + + #[tokio::test] + async fn test_calculate_data_stats_includes_blob_v2_payloads() { + const INLINE_PAYLOAD: usize = 100; + const PACKED_PAYLOAD: usize = 10 * 1024; + const DEDICATED_PAYLOAD: usize = 100 * 1024; + const PAYLOAD_TOTAL: u64 = (INLINE_PAYLOAD + PACKED_PAYLOAD + DEDICATED_PAYLOAD) as u64; + + // Thresholds picked so the three payloads land in inline, packed, and + // dedicated storage respectively. + let blobs_field = blob_field_with_options( + "blobs", + true, + BlobFieldOptions::default() + .with_inline_size_threshold(1024) + .with_dedicated_size_threshold(NonZeroUsize::new(64 * 1024).unwrap()), + ); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("i", DataType::Int32, false), + blobs_field, + ])); + let mut blobs = BlobArrayBuilder::new(3); + blobs.push_bytes(vec![1u8; INLINE_PAYLOAD]).unwrap(); + blobs.push_bytes(vec![2u8; PACKED_PAYLOAD]).unwrap(); + blobs.push_bytes(vec![3u8; DEDICATED_PAYLOAD]).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + blobs.finish().unwrap(), + ], + ) + .unwrap(); + + let test_dir = TempStrDir::default(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let dataset = Dataset::write( + reader, + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(); + + // The writer must have recorded the payload bytes in the manifest, + // attributed to the blob field. + let blob_field_id = dataset.schema().field("blobs").unwrap().id; + let fragments = dataset.fragments().clone(); + let data_file = &fragments[0].files[0]; + assert_eq!(data_file.blob_bytes.len(), data_file.fields.len()); + let recorded: u64 = data_file + .fields + .iter() + .zip(data_file.blob_bytes.iter()) + .filter(|(field_id, _)| **field_id == blob_field_id) + .map(|(_, bytes)| *bytes) + .sum(); + assert_eq!(recorded, PAYLOAD_TOTAL); + + let stats = Arc::new(dataset).calculate_data_stats().await.unwrap(); + let blob_stats = stats + .fields + .iter() + .find(|f| f.id == blob_field_id as u32) + .unwrap(); + assert!( + blob_stats.bytes_on_disk >= PAYLOAD_TOTAL, + "blob field bytes_on_disk ({}) should include the {} payload bytes", + blob_stats.bytes_on_disk, + PAYLOAD_TOTAL + ); + // The only addition beyond the payloads is the small descriptor + // column, so a tight upper bound catches double counting. + assert!( + blob_stats.bytes_on_disk < PAYLOAD_TOTAL + 4096, + "blob field bytes_on_disk ({}) suggests payloads were double counted", + blob_stats.bytes_on_disk + ); + } + + #[tokio::test] + async fn test_blob_v2_payload_bytes_survive_compaction() { + const PAYLOAD: usize = 10 * 1024; + const NUM_FRAGMENTS: u64 = 2; + const PAYLOAD_TOTAL: u64 = NUM_FRAGMENTS * PAYLOAD as u64; + + let blobs_field = blob_field_with_options( + "blobs", + true, + BlobFieldOptions::default().with_inline_size_threshold(1024), + ); + let schema = Arc::new(ArrowSchema::new(vec![blobs_field])); + + let test_dir = TempStrDir::default(); + let mut dataset = None; + for fragment_idx in 0..NUM_FRAGMENTS { + let mut blobs = BlobArrayBuilder::new(1); + blobs.push_bytes(vec![fragment_idx as u8; PAYLOAD]).unwrap(); + let batch = + RecordBatch::try_new(schema.clone(), vec![blobs.finish().unwrap()]).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + dataset = Some( + Dataset::write( + reader, + &test_dir, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + mode: if fragment_idx == 0 { + crate::dataset::WriteMode::Create + } else { + crate::dataset::WriteMode::Append + }, + ..Default::default() + }), + ) + .await + .unwrap(), + ); + } + let mut dataset = dataset.unwrap(); + assert_eq!(dataset.fragments().len(), NUM_FRAGMENTS as usize); + + // Compaction re-materializes blob payloads into new sidecar files; + // the rewritten data files must carry a fresh, correct tally. + crate::dataset::optimize::compact_files(&mut dataset, Default::default(), None) + .await + .unwrap(); + assert_eq!(dataset.fragments().len(), 1); + + let blob_field_id = dataset.schema().field("blobs").unwrap().id; + let stats = Arc::new(dataset).calculate_data_stats().await.unwrap(); + let blob_stats = stats + .fields + .iter() + .find(|f| f.id == blob_field_id as u32) + .unwrap(); + assert!( + blob_stats.bytes_on_disk >= PAYLOAD_TOTAL, + "blob field bytes_on_disk ({}) should include the {} payload bytes after compaction", + blob_stats.bytes_on_disk, + PAYLOAD_TOTAL + ); + assert!( + blob_stats.bytes_on_disk < PAYLOAD_TOTAL + 4096, + "blob field bytes_on_disk ({}) suggests payloads were double counted after compaction", + blob_stats.bytes_on_disk + ); + } } diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index 03c6f1b58f3..041d4624f77 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -877,6 +877,7 @@ async fn test_datafile_partial_replacement() { file_minor_version: minor, file_size_bytes: CachedFileSize::unknown(), base_id: None, + blob_bytes: Arc::from([]), }; let dataset = Dataset::commit( @@ -938,6 +939,7 @@ async fn test_datafile_partial_replacement() { file_minor_version: minor, file_size_bytes: CachedFileSize::unknown(), base_id: None, + blob_bytes: Arc::from([]), }; let dataset = Dataset::commit( @@ -1037,6 +1039,7 @@ async fn test_datafile_replacement_error() { file_minor_version: 0, file_size_bytes: CachedFileSize::unknown(), base_id: None, + blob_bytes: Arc::from([]), }; let new_data_file = DataFile { diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index 54e217f7437..23495006dde 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -2305,6 +2305,10 @@ impl Transaction { file.path = new_file.path.clone(); file.file_size_bytes = new_file.file_size_bytes.clone(); file.base_id = new_file.base_id; + // Carry the replacement file's blob tally (possibly + // empty = unknown); the old file's tally describes + // payloads that no longer back this path. + file.blob_bytes = new_file.blob_bytes.clone(); } columns_covered.extend(file.fields.iter()); } @@ -4387,6 +4391,7 @@ mod tests { file_minor_version: 0, file_size_bytes: CachedFileSize::new(1000), base_id: None, + blob_bytes: Arc::from([]), }); // Add a data file with all fields tombstoned @@ -4398,6 +4403,7 @@ mod tests { file_minor_version: 0, file_size_bytes: CachedFileSize::new(500), base_id: None, + blob_bytes: Arc::from([]), }); // Add a data file with mixed tombstoned and valid fields @@ -4409,6 +4415,7 @@ mod tests { file_minor_version: 0, file_size_bytes: CachedFileSize::new(750), base_id: None, + blob_bytes: Arc::from([]), }); // Add another fully tombstoned file @@ -4420,6 +4427,7 @@ mod tests { file_minor_version: 0, file_size_bytes: CachedFileSize::new(250), base_id: None, + blob_bytes: Arc::from([]), }); let mut fragments = vec![fragment]; diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index d87ccbd1662..8d477eabdc0 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -1552,6 +1552,16 @@ impl GenericWriter for V2WriterAdapter { .collect::>(); let file_version = ConcreteFileVersion::from(self.writer.version()); let write_summary = self.writer.finish().await?; + // Record the preprocessor's blob payload tally, aligned to `fields`. + // Left empty ("not recorded") when the schema has no blob v2 fields or + // when the preprocessor could not attribute every payload byte. + let blob_bytes = match self.preprocessor.as_ref().and_then(|pre| pre.blob_bytes()) { + Some(tally) if !tally.is_empty() => field_ids + .iter() + .map(|field_id| tally.get(field_id).copied().unwrap_or(0)) + .collect(), + _ => Vec::new(), + }; let data_file = DataFile::new( std::mem::take(&mut self.path), field_ids, @@ -1559,7 +1569,8 @@ impl GenericWriter for V2WriterAdapter { file_version, NonZero::new(write_summary.size_bytes), self.base_id, - ); + ) + .with_blob_bytes(blob_bytes); Ok((write_summary.num_rows as u32, data_file)) } diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index b076e778128..28ebd1ec09f 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -571,6 +571,7 @@ mod tests { file_minor_version: minor_version, file_size_bytes: CachedFileSize::new(100), base_id: None, + blob_bytes: Arc::from([]), }], overlays: vec![], deletion_file: None,