diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index 5d7a685ef..8287ef2fe 100644 --- a/python/python/lancedb/permutation.py +++ b/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 diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 2c0e2d5c1..76b3702a1 100644 --- a/python/python/lancedb/streaming.py +++ b/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: diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index 135742c84..142fc84f1 100644 --- a/python/python/tests/test_permutation.py +++ b/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( diff --git a/rust/lancedb/src/dataloader/permutation/builder.rs b/rust/lancedb/src/dataloader/permutation/builder.rs index 641c191d5..ba0d0e180 100644 --- a/rust/lancedb/src/dataloader/permutation/builder.rs +++ b/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, ) -> Result { - 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::()) + .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::() + .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::()) + .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::() + .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::()) + .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::()) + .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::()) + .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(); diff --git a/rust/lancedb/src/dataloader/permutation/reader.rs b/rust/lancedb/src/dataloader/permutation/reader.rs index 6da92e986..9757dc552 100644 --- a/rust/lancedb/src/dataloader/permutation/reader.rs +++ b/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>, split: u64, ) -> Result { + 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, + permutation_table: &Arc, + ) -> Result> { + 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::().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, permutation_table: Arc, @@ -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::()) + .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::( + 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() diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 633d0ffce..10980d55a 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/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"); diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 9feb9d5ab..9b81786cd 100644 --- a/rust/lancedb/src/table/query.rs +++ b/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);