feat: pin the base table version for data loader reads (#3982)

A permutation stores `_rowid`s, which are row addresses unless stable
row ids are enabled. Nothing in the data loader pinned a table version,
so a compaction between building a permutation and reading it can
resolve those ids to different rows.

The exposure differs by backend but exists on both:

- Remote never pins. `prepare_query_bodies` stamps `"version":
current_version()` on every request, but `current_version()` is `None`
unless `checkout` was called, so every request means "latest".
- Native pins implicitly by holding an `Arc<Dataset>` under
`ConsistencyMode::Lazy`, but `StreamingDataset.__setstate__` reopens the
table in each DataLoader worker, so each worker pins to whatever is
latest at fork time.

## Changes

`Table::at_version` returns an independent handle pinned to a version
without mutating the receiver. `checkout` cannot serve this: on remote
the version cell is an `Arc<RwLock<Option<u64>>>` shared across clones,
so pinning through it would silently pin the caller's table too.

`PermutationBuilder::build` pins for the whole build and records the
version in the permutation table's schema metadata, alongside the
existing split names. `PermutationReader` pins the base table to that
version before any take.

Because the reader pins on construction, the Python worker fork is
covered without touching the pickle format — `Permutation.__setstate__`
drops the reader and `_ensure_open` rebuilds it, which re-pins.

## Behaviour change

A permutation is now bound to the version it was built against, so rows
appended to the base table afterwards are not visible through an
existing permutation. That is the intended semantics — the permutation
only addresses rows that existed when it was built — but it is a change
worth flagging.

Permutations written before this carry no version key and read exactly
as they did before.

python/python/lancedb/permutation.py

@@ -391,6 +391,15 @@ def _table_to_pickle_state(table: Table) -> dict[str, Any]:
     }
 
 
+def _drop_base_version(permutation_data: pa.Table) -> pa.Table:
+    """Strip the recorded base version so the reader leaves the base table unpinned."""
+    metadata = dict(permutation_data.schema.metadata or {})
+    if metadata.pop(b"base_version", None) is None:
+        return permutation_data
+    metadata.pop(b"base_branch", None)
+    return permutation_data.replace_schema_metadata(metadata)
+
+
 def _table_from_pickle_state(state: dict[str, Any]) -> Table:
     from . import connect
 
@@ -679,11 +688,15 @@ class Permutation:
         from . import connect
 
         connection_factory = state["connection_factory"]
+        rebuilt_base = False
         if connection_factory is not None:
             base_table = connection_factory(state["base_table_name"])
         elif "base_table_state" in state:
-            base_table = _table_from_pickle_state(state["base_table_state"])
+            base_state = state["base_table_state"]
+            rebuilt_base = base_state["kind"] == "memory"
+            base_table = _table_from_pickle_state(base_state)
         elif "base_table_data" in state:
+            rebuilt_base = True
             # In-memory base table inlined into the pickle; rebuild the same
             # way we rebuild the in-memory permutation table.
             mem_db = connect("memory://")
@@ -701,11 +714,14 @@ class Permutation:
             )
 
         permutation_table: Optional[Table] = None
-        if state["permutation_data"] is not None:
+        permutation_data = state["permutation_data"]
+        if permutation_data is not None:
+            if rebuilt_base:
+                # The base table was materialized from Arrow, so it is a fresh
+                # single-version dataset and the recorded pin cannot resolve on it.
+                permutation_data = _drop_base_version(permutation_data)
             mem_db = connect("memory://")
-            permutation_table = mem_db.create_table(
-                "permutation", state["permutation_data"]
-            )
+            permutation_table = mem_db.create_table("permutation", permutation_data)
 
         self.base_table = base_table
         self.permutation_table = permutation_table

python/python/lancedb/streaming.py

@@ -41,6 +41,7 @@ from .permutation import (
     Permutation,
     Transforms,
     permutation_builder,
+    _drop_base_version,
     _table_from_pickle_state,
     _table_to_pickle_state,
 )
@@ -1327,6 +1328,9 @@ class StreamingDataset(IterableDataset):
             self._table = self._connection_factory(table_name)
         else:
             self._table = _table_from_pickle_state(table_state)
+            if table_state["kind"] == "memory":
+                # Rebuilt from Arrow, so the recorded pin cannot resolve on it.
+                perm_data = _drop_base_version(perm_data)
         self._perm_table = _connect("memory://").create_table(perm_name, perm_data)
 
     def state_dict(self) -> dict:

python/python/tests/test_permutation.py

@@ -56,6 +56,31 @@ def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
     assert permutation_tbl._conn.read_consistency_interval is None
 
 
+def test_pickled_permutation_reads_pinned_version(tmp_path):
+    """An unpickled copy must still read the pinned version, which also covers the
+    version surviving the ``to_arrow()`` round trip in ``__getstate__``."""
+    import pickle
+
+    db = connect(tmp_path)
+    tbl = db.create_table("base", pa.table({"idx": range(20)}))
+    permutation_tbl = permutation_builder(tbl).execute()
+    perm = Permutation.from_tables(tbl, permutation_tbl)
+
+    payload = pickle.dumps(perm)
+
+    # Compact so the stored row addresses no longer describe these rows at latest.
+    tbl.delete("true")
+    tbl.optimize()
+    assert tbl.count_rows() == 0
+
+    # Unpickle after the mutation: __setstate__ reopens at latest, so this only
+    # passes if the recorded version is applied on reopen.
+    restored = pickle.loads(payload)
+    assert len(restored) == 20
+    rows = restored.__getitems__(list(range(20)))
+    assert sorted(row["idx"] for row in rows) == list(range(20))
+
+
 def test_split_random_counts(mem_db):
     """Test random splitting with absolute counts."""
     tbl = mem_db.create_table(

rust/lancedb/src/dataloader/permutation/builder.rs

@@ -27,6 +27,12 @@ pub const SRC_ROW_ID_COL: &str = "row_id";
 
 pub const SPLIT_NAMES_CONFIG_KEY: &str = "split_names";
 
+/// Base table version the permutation was built against.
+pub const BASE_VERSION_CONFIG_KEY: &str = "base_version";
+
+/// Base table branch the permutation was built against.  Absent means main.
+pub const BASE_BRANCH_CONFIG_KEY: &str = "base_branch";
+
 pub const DEFAULT_MEMORY_LIMIT: usize = 100 * 1024 * 1024;
 
 /// Where to store the permutation table
@@ -214,21 +220,11 @@ impl PermutationBuilder {
         Ok(Box::pin(SimpleRecordBatchStream { schema, stream }))
     }
 
-    fn add_split_names(
+    fn add_config_metadata(
         data: SendableRecordBatchStream,
-        split_names: &[String],
+        metadata: HashMap<String, String>,
     ) -> Result<SendableRecordBatchStream> {
-        let schema = data
-            .schema()
-            .as_ref()
-            .clone()
-            .with_metadata(HashMap::from([(
-                SPLIT_NAMES_CONFIG_KEY.to_string(),
-                serde_json::to_string(split_names).map_err(|e| Error::Other {
-                    message: format!("Failed to serialize split names: {}", e),
-                    source: Some(e.into()),
-                })?,
-            )]));
+        let schema = data.schema().as_ref().clone().with_metadata(metadata);
         let schema = Arc::new(schema);
         let schema_clone = schema.clone();
         let stream = data.map_ok(move |batch| batch.with_schema(schema.clone()).unwrap());
@@ -269,6 +265,12 @@ impl PermutationBuilder {
             Err(err) => return Err(err),
         }
 
+        // The handle above is already pinned to one version.  Record which one, so a
+        // reader -- in a DataLoader worker, against a table that has since moved --
+        // resolves these row addresses against the same snapshot.
+        let base_version = self.base_table.version().await?;
+        let base_branch = self.base_table.current_branch();
+
         // First pass, apply filter and load row ids.  `Shuffler` permutes positions, so
         // every rank must scan the rows in the same order to build the same permutation.
         let mut rows = self.base_table.query().select(Select::columns(&[ROW_ID]));
@@ -330,11 +332,24 @@ impl PermutationBuilder {
         // Rename _rowid to row_id
         let renamed = rename_column(sorted, ROW_ID, SRC_ROW_ID_COL)?;
 
-        let streaming_data = if let Some(split_names) = &self.config.split_names {
-            Self::add_split_names(renamed, split_names)?
-        } else {
-            renamed
-        };
+        let mut metadata = HashMap::from([(
+            BASE_VERSION_CONFIG_KEY.to_string(),
+            base_version.to_string(),
+        )]);
+        // Version numbers are per-branch, so the branch is part of the coordinate.
+        if let Some(branch) = &base_branch {
+            metadata.insert(BASE_BRANCH_CONFIG_KEY.to_string(), branch.clone());
+        }
+        if let Some(split_names) = &self.config.split_names {
+            metadata.insert(
+                SPLIT_NAMES_CONFIG_KEY.to_string(),
+                serde_json::to_string(split_names).map_err(|e| Error::Other {
+                    message: format!("Failed to serialize split names: {}", e),
+                    source: Some(e.into()),
+                })?,
+            );
+        }
+        let streaming_data = Self::add_config_metadata(renamed, metadata)?;
 
         let (name, database) = match &self.config.destination {
             PermutationDestination::Permanent(database, table_name) => {
@@ -533,6 +548,141 @@ mod tests {
         assert_eq!(*planning_versions.lock().unwrap(), vec![7, 7, 6, 6]);
     }
 
+    #[tokio::test]
+    async fn test_permutation_records_base_version() {
+        let temp_dir = tempfile::tempdir().unwrap();
+
+        let db = connect(temp_dir.path().to_str().unwrap())
+            .execute()
+            .await
+            .unwrap();
+
+        let initial_data = lance_datagen::gen_batch()
+            .col("col_a", lance_datagen::array::step::<Int32Type>())
+            .into_ldb_stream(RowCount::from(100), BatchCount::from(2));
+        let data_table = db
+            .create_table("base_tbl", initial_data)
+            .execute()
+            .await
+            .unwrap();
+
+        let build_version = data_table.version().await.unwrap();
+        let permutation_table = PermutationBuilder::new(data_table.clone())
+            .build()
+            .await
+            .unwrap();
+
+        let recorded = permutation_table
+            .schema()
+            .await
+            .unwrap()
+            .metadata
+            .get(BASE_VERSION_CONFIG_KEY)
+            .expect("permutation should record the base version")
+            .parse::<u64>()
+            .unwrap();
+        assert_eq!(recorded, build_version);
+
+        // Advancing the base table must not move the recorded version.
+        let more_data = lance_datagen::gen_batch()
+            .col("col_a", lance_datagen::array::step::<Int32Type>())
+            .into_ldb_stream(RowCount::from(50), BatchCount::from(1));
+        data_table.add(more_data).execute().await.unwrap();
+        assert!(data_table.version().await.unwrap() > recorded);
+        assert_eq!(
+            permutation_table
+                .schema()
+                .await
+                .unwrap()
+                .metadata
+                .get(BASE_VERSION_CONFIG_KEY)
+                .unwrap()
+                .parse::<u64>()
+                .unwrap(),
+            recorded,
+        );
+    }
+
+    /// Version numbers are per-branch, so a permutation built on a branch must record
+    /// it -- a worker reopens by name and lands on main at the same number.
+    #[tokio::test]
+    async fn test_permutation_records_base_branch() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let db = connect(temp_dir.path().to_str().unwrap())
+            .execute()
+            .await
+            .unwrap();
+
+        let initial_data = lance_datagen::gen_batch()
+            .col("col_a", lance_datagen::array::step::<Int32Type>())
+            .into_ldb_stream(RowCount::from(10), BatchCount::from(1));
+        let data_table = db
+            .create_table("base_tbl", initial_data)
+            .execute()
+            .await
+            .unwrap();
+
+        let branch = data_table
+            .create_branch("exp", lance::dataset::refs::Ref::from(("main", 1)))
+            .await
+            .unwrap();
+        let permutation_table = PermutationBuilder::new(branch.clone())
+            .build()
+            .await
+            .unwrap();
+
+        let metadata = permutation_table.schema().await.unwrap().metadata.clone();
+        assert_eq!(
+            metadata.get(BASE_BRANCH_CONFIG_KEY).map(String::as_str),
+            Some("exp")
+        );
+
+        // Main records nothing, so an absent key keeps meaning main.
+        let main_permutation = PermutationBuilder::new(data_table.clone())
+            .build()
+            .await
+            .unwrap();
+        assert!(
+            !main_permutation
+                .schema()
+                .await
+                .unwrap()
+                .metadata
+                .contains_key(BASE_BRANCH_CONFIG_KEY)
+        );
+    }
+
+    #[tokio::test]
+    async fn test_build_does_not_pin_the_callers_table() {
+        let temp_dir = tempfile::tempdir().unwrap();
+
+        let db = connect(temp_dir.path().to_str().unwrap())
+            .execute()
+            .await
+            .unwrap();
+
+        let initial_data = lance_datagen::gen_batch()
+            .col("col_a", lance_datagen::array::step::<Int32Type>())
+            .into_ldb_stream(RowCount::from(100), BatchCount::from(1));
+        let data_table = db
+            .create_table("base_tbl", initial_data)
+            .execute()
+            .await
+            .unwrap();
+
+        PermutationBuilder::new(data_table.clone())
+            .build()
+            .await
+            .unwrap();
+
+        // The builder pins its own handle; the caller's must still track latest.
+        let more_data = lance_datagen::gen_batch()
+            .col("col_a", lance_datagen::array::step::<Int32Type>())
+            .into_ldb_stream(RowCount::from(50), BatchCount::from(1));
+        data_table.add(more_data).execute().await.unwrap();
+        assert_eq!(data_table.count_rows(None).await.unwrap(), 150);
+    }
+
     #[tokio::test]
     async fn test_permutation_builder() {
         let temp_dir = tempfile::tempdir().unwrap();

rust/lancedb/src/dataloader/permutation/reader.rs

@@ -8,7 +8,9 @@
 //! the rows from a source table that correspond to row IDs stored in a separate table.
 
 use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream};
-use crate::dataloader::permutation::builder::SRC_ROW_ID_COL;
+use crate::dataloader::permutation::builder::{
+    BASE_BRANCH_CONFIG_KEY, BASE_VERSION_CONFIG_KEY, SRC_ROW_ID_COL,
+};
 use crate::dataloader::permutation::split::SPLIT_ID_COLUMN;
 use crate::error::Error;
 use crate::query::{
@@ -23,6 +25,7 @@ use arrow_array::{RecordBatch, UInt64Array};
 use arrow_schema::SchemaRef;
 use datafusion_expr::{Expr, col, lit};
 use futures::{StreamExt, TryStreamExt};
+use lance::dataset::refs::MAIN_BRANCH;
 use lance::dataset::scanner::DatasetRecordBatchStream;
 use lance::io::RecordBatchStream;
 use lance_arrow::RecordBatchExt;
@@ -69,6 +72,10 @@ impl PermutationReader {
         permutation_table: Option<Arc<dyn BaseTable>>,
         split: u64,
     ) -> Result<Self> {
+        let base_table = match &permutation_table {
+            Some(permutation_table) => Self::pin_base_table(base_table, permutation_table).await?,
+            None => base_table,
+        };
         let mut slf = Self {
             base_table,
             permutation_table,
@@ -89,6 +96,34 @@ impl PermutationReader {
         Ok(slf)
     }
 
+    /// Pins the base table to the version the permutation was built against.
+    /// Permutations written before that was recorded carry no key and stay unpinned.
+    async fn pin_base_table(
+        base_table: Arc<dyn BaseTable>,
+        permutation_table: &Arc<dyn BaseTable>,
+    ) -> Result<Arc<dyn BaseTable>> {
+        let schema = permutation_table.schema().await?;
+        let Some(raw) = schema.metadata.get(BASE_VERSION_CONFIG_KEY) else {
+            return Ok(base_table);
+        };
+        let version = raw.parse::<u64>().map_err(|e| Error::InvalidInput {
+            message: format!(
+                "Permutation table has an unreadable {} of {:?}: {}",
+                BASE_VERSION_CONFIG_KEY, raw, e
+            ),
+        })?;
+        // The recorded branch, not the handle's: a worker reopens by name and lands
+        // on main, and version numbers are per-branch.
+        let branch = schema
+            .metadata
+            .get(BASE_BRANCH_CONFIG_KEY)
+            .map(String::as_str)
+            .unwrap_or(MAIN_BRANCH);
+        base_table
+            .checkout_branch_version(branch, Some(version))
+            .await
+    }
+
     pub async fn try_from_tables(
         base_table: Arc<dyn BaseTable>,
         permutation_table: Arc<dyn BaseTable>,
@@ -511,9 +546,13 @@ mod tests {
     use lance_datagen::{BatchCount, RowCount};
     use rand::seq::SliceRandom;
 
+    // Aliased: `test_utils::datagen` exports a trait of the same name.
+    use crate::arrow::LanceDbDatagenExt as _;
     use crate::{
         Table,
         arrow::SendableRecordBatchStream,
+        connect,
+        dataloader::permutation::builder::PermutationBuilder,
         query::{ExecutableQuery, QueryBase},
         test_utils::datagen::{LanceDbDatagenExt, virtual_table},
     };
@@ -545,6 +584,58 @@ mod tests {
         .await
     }
 
+    /// Compaction moves row addresses, so the reader must read the pinned version.
+    #[tokio::test]
+    async fn test_reader_pins_base_version() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let db = connect(temp_dir.path().to_str().unwrap())
+            .execute()
+            .await
+            .unwrap();
+
+        let data = lance_datagen::gen_batch()
+            .col("idx", lance_datagen::array::step::<Int32Type>())
+            .into_ldb_stream(RowCount::from(20), BatchCount::from(1));
+        let base_table = db.create_table("base_tbl", data).execute().await.unwrap();
+
+        let permutation_table = PermutationBuilder::new(base_table.clone())
+            .build()
+            .await
+            .unwrap();
+
+        base_table.delete("true").await.unwrap();
+        base_table
+            .optimize(crate::table::OptimizeAction::All)
+            .await
+            .unwrap();
+        assert_eq!(base_table.count_rows(None).await.unwrap(), 0);
+
+        let reader = PermutationReader::try_from_tables(
+            base_table.base_table().clone(),
+            permutation_table.base_table().clone(),
+            0,
+        )
+        .await
+        .unwrap();
+
+        let values = collect_from_stream::<Int32Type>(
+            reader
+                .read(
+                    Select::Columns(vec!["idx".to_string()]),
+                    QueryExecutionOptions::default(),
+                )
+                .await
+                .unwrap(),
+            "idx",
+        )
+        .await;
+        assert_eq!(
+            values.len(),
+            20,
+            "reader should still see the pinned version"
+        );
+    }
+
     #[tokio::test]
     async fn test_permutation_reader() {
         let base_table = lance_datagen::gen_batch()

rust/lancedb/src/remote/table.rs

@@ -4349,6 +4349,43 @@ mod tests {
         assert!(!table.base_table().scan_order_is_deterministic());
     }
 
+    #[tokio::test]
+    async fn test_checkout_branch_pins_without_touching_the_original() {
+        let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
+        let recorder = seen.clone();
+        let table = Table::new_with_handler_version(
+            "my_table",
+            semver::Version::new(0, 5, 0),
+            move |request| match request.url().path() {
+                "/v1/table/my_table/describe/" => http::Response::builder()
+                    .status(200)
+                    .body(br#"{"version": 42, "schema": {"fields": []}}"#.to_vec())
+                    .unwrap(),
+                "/v1/table/my_table/count_rows/" => {
+                    let body = request_body_json(&request);
+                    recorder.lock().unwrap().push(body["version"].clone());
+                    http::Response::builder()
+                        .status(200)
+                        .body(b"0".to_vec())
+                        .unwrap()
+                }
+                path => panic!("unexpected request path: {path}"),
+            },
+        );
+
+        let pinned = table.checkout_branch("main", Some(42)).await.unwrap();
+        pinned.count_rows(None).await.unwrap();
+        table.count_rows(None).await.unwrap();
+
+        let seen = seen.lock().unwrap();
+        assert_eq!(seen[0], 42, "the pinned handle must send its version");
+        assert!(
+            seen[1].is_null(),
+            "the original handle must still track latest, got {:?}",
+            seen[1]
+        );
+    }
+
     #[tokio::test]
     async fn test_fetch_blobs_sends_the_checked_out_version() {
         let ipc = one_row_blob_ipc_stream("image");

rust/lancedb/src/table/query.rs

@@ -70,6 +70,9 @@ async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> R
         .contains(&NamespaceClientPushdownOperation::QueryTable)
         && table.namespace_client.is_some()
         && table.dataset.current_branch().is_none()
+        // NsQueryTableRequest has no version field, so a pushed-down query would
+        // read latest and ignore the pin.
+        && table.dataset.time_travel_version().is_none()
         && !requires_local_namespace_execution(query))
     {
         return Ok(false);