Skip to content

Commit 9931ace

Browse files
beinanclaude
andauthored
perf: intern DataFile fields/column_indices to reduce manifest memory (#6477)
## Summary - Change `DataFile.fields` and `DataFile.column_indices` from `Vec<i32>` to `Arc<[i32]>` so that fragments with identical field lists share a single heap allocation - Add `DataFileFieldInterner` that deduplicates these slices during manifest deserialization - In homogeneous tables (the common case), every fragment carries the same field list, so at 20M fragments this saves **~2.4 GB** of redundant heap allocations ## Motivation When dataset manifests grow large (>1 GB with millions of fragments), opening the dataset becomes very expensive in terms of memory. Each `DataFile` previously owned its own `Vec<i32>` for `fields` and `column_indices`, even though in most tables every fragment has the exact same field list. This PR deduplicates those allocations at deserialization time. ### Per-fragment memory breakdown (before) | Field | Size per fragment | |-------|------------------| | `fields: Vec<i32>` (10 fields) | ~64 bytes | | `column_indices: Vec<i32>` (10 cols) | ~64 bytes | | **Total redundant** | **~128 bytes x 20M = ~2.4 GB** | ### After this change With interning, all 20M fragments share a single `Arc<[i32]>` allocation (~80 bytes total instead of 2.4 GB). ## Changes - **`lance-table/src/format/fragment.rs`** — Core struct change (`Vec<i32>` → `Arc<[i32]>`), custom `Serialize`/`Deserialize` impls, and `DataFileFieldInterner` - **`lance-table/src/format/manifest.rs`** — Use interner during manifest deserialization - **`lance/src/dataset/fragment.rs`**, **`merge_insert.rs`**, **`io/commit.rs`** — Tombstoning and field-remapping rebuilt as new `Arc<[i32]>` instead of in-place mutation - **`python/src/fragment.rs`**, **`java/lance-jni/src/fragment.rs`** — FFI boundary conversions - Various test files — Updated struct literals and assertions ## Compatibility - No format change — protobuf schema is unchanged - Serde JSON output is identical (custom impl serializes `Arc<[i32]>` as `[i32]`) - All public API signatures that take `Vec<i32>` (e.g., `DataFile::new()`, `Fragment::add_file()`) still accept `Vec<i32>` and convert internally 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 27351e5 commit 9931ace

17 files changed

Lines changed: 257 additions & 104 deletions

File tree

java/lance-jni/src/fragment.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -536,8 +536,8 @@ impl IntoJava for &FragmentUpdateResult {
536536
impl IntoJava for &DataFile {
537537
fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result<JObject<'a>> {
538538
let path = env.new_string(self.path.clone())?.into();
539-
let fields = JLance(self.fields.clone()).into_java(env)?;
540-
let column_indices = JLance(self.column_indices.clone()).into_java(env)?;
539+
let fields = JLance(self.fields.to_vec()).into_java(env)?;
540+
let column_indices = JLance(self.column_indices.to_vec()).into_java(env)?;
541541
let file_size_bytes = match self.file_size_bytes.get() {
542542
Some(f) => JLance(u64::from(f) as i64).into_java(env)?,
543543
None => JObject::null(),
@@ -775,8 +775,8 @@ impl FromJObjectWithEnv<DataFile> for JObject<'_> {
775775
let base_id = get_base_id(env, self)?;
776776
Ok(DataFile {
777777
path,
778-
fields,
779-
column_indices,
778+
fields: fields.into(),
779+
column_indices: column_indices.into(),
780780
file_major_version,
781781
file_minor_version,
782782
file_size_bytes,

python/src/debug.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,8 @@ impl PrettyPrintableFragment {
6161
let schema = schema.project_by_ids(&file.fields, false);
6262
PrettyPrintableDataFile {
6363
path: file.path.clone(),
64-
fields: file.fields.clone(),
65-
column_indices: file.column_indices.clone(),
64+
fields: file.fields.to_vec(),
65+
column_indices: file.column_indices.to_vec(),
6666
schema,
6767
major_version: file.file_major_version,
6868
minor_version: file.file_minor_version,

python/src/fragment.rs

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -846,10 +846,12 @@ impl FromPyObject<'_> for PyLance<DataFile> {
846846
fn extract_bound(ob: &pyo3::Bound<'_, PyAny>) -> PyResult<Self> {
847847
let file_size_bytes: Option<u64> = ob.getattr("file_size_bytes")?.extract()?;
848848
let file_size_bytes = CachedFileSize::new(file_size_bytes.unwrap_or(0));
849+
let fields: Vec<i32> = ob.getattr("fields")?.extract()?;
850+
let column_indices: Vec<i32> = ob.getattr("column_indices")?.extract()?;
849851
Ok(Self(DataFile {
850852
path: ob.getattr("path")?.extract()?,
851-
fields: ob.getattr("fields")?.extract()?,
852-
column_indices: ob.getattr("column_indices")?.extract()?,
853+
fields: fields.into(),
854+
column_indices: column_indices.into(),
853855
file_major_version: ob.getattr("file_major_version")?.extract()?,
854856
file_minor_version: ob.getattr("file_minor_version")?.extract()?,
855857
file_size_bytes,
@@ -872,8 +874,8 @@ impl<'py> IntoPyObject<'py> for PyLance<&DataFile> {
872874
let file_size_bytes = self.0.file_size_bytes.get().map(u64::from);
873875
cls.call1((
874876
&self.0.path,
875-
self.0.fields.clone(),
876-
self.0.column_indices.clone(),
877+
self.0.fields.to_vec(),
878+
self.0.column_indices.to_vec(),
877879
self.0.file_major_version,
878880
self.0.file_minor_version,
879881
file_size_bytes,

rust/lance-table/src/format/fragment.rs

Lines changed: 129 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright The Lance Authors
33

4+
use std::collections::HashMap;
45
use std::num::NonZero;
6+
use std::sync::Arc;
57

68
use deepsize::DeepSizeOf;
79
use lance_core::Error;
810
use lance_file::format::{MAJOR_VERSION, MINOR_VERSION};
911
use lance_file::version::LanceFileVersion;
1012
use lance_io::utils::CachedFileSize;
1113
use object_store::path::Path;
12-
use serde::{Deserialize, Serialize};
14+
use serde::{Deserialize, Deserializer, Serialize, Serializer};
1315

1416
use crate::format::pb;
1517

@@ -22,12 +24,16 @@ use lance_core::error::Result;
2224
/// Lance Data File
2325
///
2426
/// A data file is one piece of file storing data.
25-
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
27+
#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)]
2628
pub struct DataFile {
2729
/// Relative path of the data file to dataset root.
2830
pub path: String,
2931
/// The ids of fields in this file.
30-
pub fields: Vec<i32>,
32+
///
33+
/// When identical across many fragments (common case), multiple `DataFile`
34+
/// instances share a single heap allocation via `Arc`, significantly
35+
/// reducing manifest memory for large tables.
36+
pub fields: Arc<[i32]>,
3137
/// The offsets of the fields listed in `fields`, empty in v1 files
3238
///
3339
/// Note that -1 is a possibility and it indices that the field has
@@ -37,13 +43,10 @@ pub struct DataFile {
3743
/// `column_indices`; such columns are ignored by field-id–based projection.
3844
/// For example, some fields, such as blob fields, occupy multiple
3945
/// columns in the file but only have a single field id.
40-
#[serde(default)]
41-
pub column_indices: Vec<i32>,
46+
pub column_indices: Arc<[i32]>,
4247
/// The major version of the file format used to write this file.
43-
#[serde(default)]
4448
pub file_major_version: u32,
4549
/// The minor version of the file format used to write this file.
46-
#[serde(default)]
4750
pub file_minor_version: u32,
4851

4952
/// The size of the file in bytes, if known.
@@ -53,6 +56,52 @@ pub struct DataFile {
5356
pub base_id: Option<u32>,
5457
}
5558

59+
// Custom Serialize: convert Arc<[i32]> to slice for transparent JSON output
60+
impl Serialize for DataFile {
61+
fn serialize<S: Serializer>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> {
62+
use serde::ser::SerializeStruct;
63+
let mut s = serializer.serialize_struct("DataFile", 7)?;
64+
s.serialize_field("path", &self.path)?;
65+
s.serialize_field("fields", self.fields.as_ref())?;
66+
s.serialize_field("column_indices", self.column_indices.as_ref())?;
67+
s.serialize_field("file_major_version", &self.file_major_version)?;
68+
s.serialize_field("file_minor_version", &self.file_minor_version)?;
69+
s.serialize_field("file_size_bytes", &self.file_size_bytes)?;
70+
s.serialize_field("base_id", &self.base_id)?;
71+
s.end()
72+
}
73+
}
74+
75+
// Custom Deserialize: read Vec<i32> and convert to Arc<[i32]>
76+
impl<'de> Deserialize<'de> for DataFile {
77+
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
78+
#[derive(Deserialize)]
79+
struct DataFileHelper {
80+
path: String,
81+
fields: Vec<i32>,
82+
#[serde(default)]
83+
column_indices: Vec<i32>,
84+
#[serde(default)]
85+
file_major_version: u32,
86+
#[serde(default)]
87+
file_minor_version: u32,
88+
file_size_bytes: CachedFileSize,
89+
base_id: Option<u32>,
90+
}
91+
92+
let helper = DataFileHelper::deserialize(deserializer)?;
93+
Ok(Self {
94+
path: helper.path,
95+
fields: Arc::from(helper.fields),
96+
column_indices: Arc::from(helper.column_indices),
97+
file_major_version: helper.file_major_version,
98+
file_minor_version: helper.file_minor_version,
99+
file_size_bytes: helper.file_size_bytes,
100+
base_id: helper.base_id,
101+
})
102+
}
103+
}
104+
56105
impl DataFile {
57106
pub fn new(
58107
path: impl Into<String>,
@@ -65,8 +114,8 @@ impl DataFile {
65114
) -> Self {
66115
Self {
67116
path: path.into(),
68-
fields,
69-
column_indices,
117+
fields: Arc::from(fields),
118+
column_indices: Arc::from(column_indices),
70119
file_major_version,
71120
file_minor_version,
72121
file_size_bytes: file_size_bytes.into(),
@@ -82,8 +131,8 @@ impl DataFile {
82131
) -> Self {
83132
Self {
84133
path: path.into(),
85-
fields: vec![],
86-
column_indices: vec![],
134+
fields: Arc::from([]),
135+
column_indices: Arc::from([]),
87136
file_major_version,
88137
file_minor_version,
89138
file_size_bytes: Default::default(),
@@ -158,8 +207,8 @@ impl From<&DataFile> for pb::DataFile {
158207
fn from(df: &DataFile) -> Self {
159208
Self {
160209
path: df.path.clone(),
161-
fields: df.fields.clone(),
162-
column_indices: df.column_indices.clone(),
210+
fields: df.fields.to_vec(),
211+
column_indices: df.column_indices.to_vec(),
163212
file_major_version: df.file_major_version,
164213
file_minor_version: df.file_minor_version,
165214
file_size_bytes: df.file_size_bytes.get().map_or(0, |v| v.get()),
@@ -174,14 +223,77 @@ impl TryFrom<pb::DataFile> for DataFile {
174223
fn try_from(proto: pb::DataFile) -> Result<Self> {
175224
Ok(Self {
176225
path: proto.path,
177-
fields: proto.fields,
178-
column_indices: proto.column_indices,
226+
fields: Arc::from(proto.fields),
227+
column_indices: Arc::from(proto.column_indices),
228+
file_major_version: proto.file_major_version,
229+
file_minor_version: proto.file_minor_version,
230+
file_size_bytes: CachedFileSize::new(proto.file_size_bytes),
231+
base_id: proto.base_id,
232+
})
233+
}
234+
}
235+
236+
/// Interns `fields` and `column_indices` slices so that `DataFile` instances
237+
/// with identical content share a single `Arc<[i32]>` allocation.
238+
///
239+
/// At 20M fragments the deduplication typically saves multiple GB of heap
240+
/// because every fragment in a homogeneous table carries the same field list.
241+
#[derive(Default)]
242+
pub struct DataFileFieldInterner {
243+
fields: HashMap<Vec<i32>, Arc<[i32]>>,
244+
column_indices: HashMap<Vec<i32>, Arc<[i32]>>,
245+
}
246+
247+
impl DataFileFieldInterner {
248+
/// Intern a `Vec<i32>`: if the same content was seen before, return a
249+
/// clone of the existing `Arc`; otherwise store and return a new one.
250+
fn intern(cache: &mut HashMap<Vec<i32>, Arc<[i32]>>, v: Vec<i32>) -> Arc<[i32]> {
251+
cache
252+
.entry(v)
253+
.or_insert_with_key(|k| Arc::from(k.as_slice()))
254+
.clone()
255+
}
256+
257+
/// Convert a protobuf `DataFile`, interning `fields` and `column_indices`.
258+
pub fn intern_data_file(&mut self, proto: pb::DataFile) -> Result<DataFile> {
259+
Ok(DataFile {
260+
path: proto.path,
261+
fields: Self::intern(&mut self.fields, proto.fields),
262+
column_indices: Self::intern(&mut self.column_indices, proto.column_indices),
179263
file_major_version: proto.file_major_version,
180264
file_minor_version: proto.file_minor_version,
181265
file_size_bytes: CachedFileSize::new(proto.file_size_bytes),
182266
base_id: proto.base_id,
183267
})
184268
}
269+
270+
/// Convert a protobuf `DataFragment`, interning fields within its data files.
271+
pub fn intern_fragment(&mut self, p: pb::DataFragment) -> Result<Fragment> {
272+
let physical_rows = if p.physical_rows > 0 {
273+
Some(p.physical_rows as usize)
274+
} else {
275+
None
276+
};
277+
Ok(Fragment {
278+
id: p.id,
279+
files: p
280+
.files
281+
.into_iter()
282+
.map(|f| self.intern_data_file(f))
283+
.collect::<Result<_>>()?,
284+
deletion_file: p.deletion_file.map(DeletionFile::try_from).transpose()?,
285+
row_id_meta: p.row_id_sequence.map(RowIdMeta::try_from).transpose()?,
286+
physical_rows,
287+
last_updated_at_version_meta: p
288+
.last_updated_at_version_sequence
289+
.map(RowDatasetVersionMeta::try_from)
290+
.transpose()?,
291+
created_at_version_meta: p
292+
.created_at_version_sequence
293+
.map(RowDatasetVersionMeta::try_from)
294+
.transpose()?,
295+
})
296+
}
185297
}
186298

187299
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
@@ -624,9 +736,9 @@ mod tests {
624736
fn data_file_validate_allows_extra_columns() {
625737
let data_file = DataFile {
626738
path: "foo.lance".to_string(),
627-
fields: vec![1, 2],
739+
fields: Arc::from([1, 2]),
628740
// One extra column without a field id mapping
629-
column_indices: vec![0, 1, 2],
741+
column_indices: Arc::from([0, 1, 2]),
630742
file_major_version: MAJOR_VERSION as u32,
631743
file_minor_version: MINOR_VERSION as u32,
632744
file_size_bytes: Default::default(),

rust/lance-table/src/format/manifest.rs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use std::sync::Arc;
1717

1818
use super::Fragment;
1919
use crate::feature_flags::{FLAG_STABLE_ROW_IDS, has_deprecated_v2_feature_flag};
20+
use crate::format::fragment::DataFileFieldInterner;
2021
use crate::format::pb;
2122
use lance_core::cache::LanceCache;
2223
use lance_core::datatypes::Schema;
@@ -426,7 +427,7 @@ impl Manifest {
426427
let fragment_max_id = self
427428
.fragments
428429
.iter()
429-
.flat_map(|f| f.files.iter().flat_map(|file| file.fields.as_slice()))
430+
.flat_map(|f| f.files.iter().flat_map(|file| file.fields.iter()))
430431
.max()
431432
.copied();
432433
let fragment_max_id = fragment_max_id.unwrap_or(-1);
@@ -863,10 +864,11 @@ impl TryFrom<pb::Manifest> for Manifest {
863864
}),
864865
_ => None,
865866
};
867+
let mut interner = DataFileFieldInterner::default();
866868
let fragments = Arc::new(
867869
p.fragments
868870
.into_iter()
869-
.map(Fragment::try_from)
871+
.map(|f| interner.intern_fragment(f))
870872
.collect::<Result<Vec<_>>>()?,
871873
);
872874
let fragment_offsets = compute_fragment_offsets(fragments.as_slice());

rust/lance/src/dataset/fragment.rs

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1209,7 +1209,7 @@ impl FileFragment {
12091209
let mut seen_fields = HashSet::new();
12101210
for data_file in &self.metadata.files {
12111211
let last = -1;
1212-
for field_id in &data_file.fields {
1212+
for field_id in data_file.fields.iter() {
12131213
if *field_id <= last {
12141214
return Err(Error::corrupt_file(
12151215
self.dataset
@@ -1666,12 +1666,19 @@ impl FileFragment {
16661666
// Mark fields in updated data files as obsolete ("tombstone").
16671667
let updated_fields = updated_fragment.files.last().unwrap().fields.clone();
16681668
for data_file in &mut updated_fragment.files.iter_mut().rev().skip(1) {
1669-
for field in &mut data_file.fields {
1670-
if updated_fields.contains(field) {
1671-
// Tombstone these fields
1672-
*field = -2;
1673-
}
1674-
}
1669+
let new_fields: Arc<[i32]> = data_file
1670+
.fields
1671+
.iter()
1672+
.map(|field| {
1673+
if updated_fields.contains(field) {
1674+
-2 // Tombstone
1675+
} else {
1676+
*field
1677+
}
1678+
})
1679+
.collect::<Vec<_>>()
1680+
.into();
1681+
data_file.fields = new_fields;
16751682
}
16761683
// Remove data files that have become entirely tombstoned.
16771684
updated_fragment

0 commit comments

Comments
 (0)