diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index bcf84bf3a..5d7a685ef 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -41,21 +41,15 @@ class PermutationBuilder: The permutation is stored in memory and will be lost when the program exits. """ - def __init__(self, table: LanceTable): + def __init__(self, table: Table): """ Creates a new permutation builder for the given table. By default, the permutation builder will create a single split that contains all rows in the same order as the base table. + + Tables with an LSM write spec are rejected: unflushed rows have no row id. """ - if not hasattr(table, "_inner"): - raise TypeError( - f"PermutationBuilder requires a local LanceTable, " - f"got {type(table).__name__}. " - "The permutation API is not supported on remote tables. " - "Remote tables connect to LanceDB Cloud or Enterprise and do not have " - "direct access to the underlying Lance dataset needed for permutations." - ) self._async = async_permutation_builder(table) def split_random( @@ -231,7 +225,7 @@ class PermutationBuilder: return LOOP.run(do_execute()) -def permutation_builder(table: LanceTable) -> PermutationBuilder: +def permutation_builder(table: Table) -> PermutationBuilder: return PermutationBuilder(table) @@ -248,7 +242,7 @@ class Permutations: Attributes ---------- - base_table: LanceTable + base_table: Table The base table that the permutations are based on. permutation_table: LanceTable The permutation table that defines the splits. @@ -282,7 +276,7 @@ class Permutations: {'train': 0, 'test': 1} """ - def __init__(self, base_table: LanceTable, permutation_table: LanceTable): + def __init__(self, base_table: Table, permutation_table: LanceTable): self.base_table = base_table self.permutation_table = permutation_table diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 734f835c6..0c1a70765 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -37,6 +37,11 @@ from unittest.mock import patch import lancedb import pyarrow as pa import pytest +from utils import ( + MockPermutationServer, + assert_server_safe_row_id_requests, + mock_remote_table, +) torch = pytest.importorskip("torch") streaming = pytest.importorskip("lancedb.streaming") @@ -2118,3 +2123,27 @@ def test_doc_example_checkpoint(lance_table): assert sorted(consumed + remaining_original) == list(range(NUM_ROWS)), ( "Consumed + remaining must cover every row exactly once" ) + + +# --------------------------------------------------------------------------- +# Remote tables (LanceDB Cloud / Enterprise) +# --------------------------------------------------------------------------- + + +def test_streaming_dataset_over_remote_table(): + """StreamingDataset reads a remote table, with server-safe requests. + + Builds a permutation over a remote table, then fetches batches from it by row id. + """ + server = MockPermutationServer() + + with mock_remote_table(server) as table: + ds = StreamingDataset(table, num_splits=2, shuffle_seed=SHUFFLE_SEED) + ids = [row["id"] for row in ds] + + assert sorted(ids) == list(range(server.num_rows)), ( + "Every row of the remote table must be yielded exactly once" + ) + assert len(server.scans) == 1, "the permutation is built with one row-id scan" + assert server.takes, "rows must be fetched with row-id takes" + assert_server_safe_row_id_requests(server) diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index 6d8f6f431..135742c84 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -8,6 +8,11 @@ import pytest from lancedb import DBConnection, Table, connect from lancedb.background_loop import LOOP from lancedb.permutation import Permutation, Permutations, permutation_builder +from utils import ( + MockPermutationServer, + assert_server_safe_row_id_requests, + mock_remote_table, +) def test_split_random_ratios(mem_db): @@ -1214,3 +1219,57 @@ def test_remove_rowid_after_select(some_permutation: Permutation): perm_without_rowid = perm_with_rowid.remove_columns(["_rowid"]) assert "_rowid" not in perm_without_rowid.column_names assert perm_without_rowid.column_names == ["id"] + + +def test_permutation_is_stable_when_remote_scan_order_varies(): + """Splits are assigned by scan position, and every rank builds its own + permutation, so two ranks seeing different scan orders must still agree.""" + server = MockPermutationServer(num_rows=16, vary_scan_order=True) + + def split_of_each_row(permutation_tbl): + # Sequential splits are assigned by position, so a reversed scan would put + # the last rows in split 0. Compare the mapping rather than the table order, + # which the split-id sort does not pin down. + rows = permutation_tbl.search(None).to_arrow().to_pydict() + return dict(zip(rows["row_id"], rows["split_id"])) + + with mock_remote_table(server) as table: + first = split_of_each_row( + permutation_builder(table).split_sequential(fixed=2).execute() + ) + second = split_of_each_row( + permutation_builder(table).split_sequential(fixed=2).execute() + ) + + assert server.scan_calls == 2, "both builds must have scanned" + assert first == second + assert first[0] == 0 and first[server.num_rows - 1] == 1, first + + +def test_permutation_over_remote_table(): + """The permutation API accepts a remote table, addressing rows by `_rowid` just + as `take_row_ids` does. Also pins the request shapes sent to the server. + """ + server = MockPermutationServer() + + with mock_remote_table(server) as table: + permutation_tbl = permutation_builder(table).split_sequential(fixed=2).execute() + assert permutation_tbl.count_rows() == server.num_rows + + permutation = Permutation.from_tables(table, permutation_tbl, 0) + assert permutation.num_rows == server.num_rows // 2 + + # Compare against the permutation's own order; the split-id sort is not stable. + rows = permutation_tbl.search(None).to_arrow().to_pydict() + split0 = [ + row_id + for row_id, split in zip(rows["row_id"], rows["split_id"]) + if not split + ] + # The mock table's `id` equals its `_rowid`. + assert permutation.take_offsets([2, 0]) == [ + {"id": split0[2]}, + {"id": split0[0]}, + ] + + assert_server_safe_row_id_requests(server) diff --git a/python/python/tests/utils.py b/python/python/tests/utils.py index 62ec74497..4882f0c28 100644 --- a/python/python/tests/utils.py +++ b/python/python/tests/utils.py @@ -1,7 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import contextlib +import http.server +import json +import re +import threading + +import lancedb +import pyarrow as pa import pytest +ARROW_FILE_CONTENT_TYPE = "application/vnd.apache.arrow.file" + def exception_output(e_info: pytest.ExceptionInfo): import traceback @@ -9,3 +19,199 @@ def exception_output(e_info: pytest.ExceptionInfo): # skip traceback part, since it's not worth checking in tests lines = traceback.format_exception_only(e_info.type, e_info.value) return "".join(lines).strip() + + +def parse_in_list(filter_sql: str) -> list[int]: + """Pull the integers out of a ` IN (a, b, c)` predicate. + + Scoped to the parenthesised list so a cast in the SQL adds no phantom values. + """ + match = re.search(r"\bIN\s*\(([^)]*)\)", filter_sql, re.IGNORECASE) + assert match is not None, f"expected an IN list, got: {filter_sql}" + return [int(m) for m in re.findall(r"-?\d+", match.group(1))] + + +def is_row_id_take(body) -> bool: + """True when a query body fetches specific rows by row id.""" + return "_rowid" in (body.get("filter") or "") + + +def arrow_file_bytes(table: pa.Table) -> bytes: + """Serialize to the Arrow IPC *file* framing the /query/ route answers with.""" + sink = pa.BufferOutputStream() + with pa.ipc.new_file(sink, table.schema) as writer: + writer.write_table(table) + return sink.getvalue().to_pybytes() + + +class MockPermutationServer: + """A stand-in LanceDB server hosting one table whose ``id`` equals its ``_rowid``. + + Records every ``/query/`` body so tests can assert on the request shapes sent to + the server, which is the part that has to stay compatible. + """ + + def __init__(self, name="remote_data", num_rows=8, vary_scan_order=False): + self.name = name + self.num_rows = num_rows + self.query_bodies = [] + # Stand in for a distributed scan that answers in no fixed order. + self.vary_scan_order = vary_scan_order + self.scan_calls = 0 + + def __call__(self, request): + path = request.path + if path == f"/v1/table/{self.name}/describe/": + return self._json( + request, + { + "version": 1, + "schema": { + "fields": [ + {"name": "id", "type": {"type": "int64"}, "nullable": False} + ] + }, + }, + ) + if path == f"/v1/table/{self.name}/get_lsm_write_spec/": + self._read_body(request) + # Null spec: this table has no LSM write path. + return self._json(request, {"lsm_write_spec": None}) + if path == f"/v1/table/{self.name}/count_rows/": + self._read_body(request) + return self._json(request, self.num_rows) + if path == f"/v1/table/{self.name}/query/": + return self._query(request, self._read_body(request)) + + # Drain first, so an unexpected route cannot desync a keep-alive connection. + self._read_body(request) + request.send_response(404) + request.end_headers() + + @property + def scans(self): + """Bodies of the permutation build scan: the row id column, nothing else.""" + return [b for b in self.query_bodies if b.get("columns") == ["_rowid"]] + + @property + def takes(self): + """Bodies of the row-id takes the loader fetches batches with. + + Keyed on `_rowid`, not "has a filter": the schema probe also has a predicate. + """ + return [b for b in self.query_bodies if is_row_id_take(b)] + + @staticmethod + def _read_body(request): + content_len = int(request.headers.get("Content-Length") or 0) + return json.loads(request.rfile.read(content_len)) if content_len else {} + + @staticmethod + def _json(request, payload): + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(payload).encode()) + + @staticmethod + def _arrow(request, table): + body = arrow_file_bytes(table) + request.send_response(200) + request.send_header("Content-Type", ARROW_FILE_CONTENT_TYPE) + request.send_header("Content-Length", str(len(body))) + request.end_headers() + request.wfile.write(body) + + def _query(self, request, body): + self.query_bodies.append(body) + + if is_row_id_take(body): + # A row-id take. Answer ascending, so tests prove the client reorders. + row_ids = sorted(parse_in_list(body["filter"])) + return self._arrow( + request, + pa.table( + { + "id": pa.array(row_ids, pa.int64()), + "_rowid": pa.array(row_ids, pa.uint64()), + } + ), + ) + + if body.get("columns") == ["_rowid"]: + # The permutation build scan: row ids and nothing else. + row_ids = list(range(self.num_rows)) + if self.vary_scan_order and self.scan_calls % 2: + row_ids.reverse() + self.scan_calls += 1 + return self._arrow( + request, + pa.table({"_rowid": pa.array(row_ids, pa.uint64())}), + ) + + # The schema probe: filtered to nothing, so it carries schema and no rows. + return self._arrow(request, pa.table({"id": pa.array([], pa.int64())})) + + +def _make_handler(serve): + class MockLanceDBHandler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + serve(self) + + def do_POST(self): + serve(self) + + def log_message(self, *args): + pass # keep pytest output readable + + return MockLanceDBHandler + + +@contextlib.contextmanager +def mock_remote_table(server): + """Run ``server`` on a local port and yield an open remote table against it. + + Threading: the loader fans out fetch threads a single-threaded server would + serialize, hiding the prefetch overlap under test. + """ + with http.server.ThreadingHTTPServer( + ("localhost", 0), _make_handler(server) + ) as srv: + thread = threading.Thread(target=srv.serve_forever) + thread.start() + try: + db = lancedb.connect( + "db://dev", + api_key="fake", + host_override=f"http://localhost:{srv.server_address[1]}", + client_config={"timeout_config": {"connect_timeout": 5}}, + ) + yield db.open_table(server.name) + finally: + srv.shutdown() + thread.join() + + +def assert_server_safe_row_id_requests(server): + """Assert the loader fetched rows by row id and bounded everything else. + + `.get`, not `[...]`, so a dropped field reads as the assertion, not a KeyError. + """ + for body in server.takes: + # The fetch needs the row id back to restore the requested order. + assert body.get("with_row_id") is True, body + assert "_rowid" in body["filter"], body + + # Only the one-off permutation scan may scan the whole table; the schema probe is + # built once per split per epoch. `k == 0` counts as unbounded: lance reads a zero + # limit as "no limit". + def is_unbounded(body): + if is_row_id_take(body): + return False + k = body.get("k") + return k is None or k == 0 or k > server.num_rows + + unbounded = [b for b in server.query_bodies if is_unbounded(b)] + assert unbounded == server.scans, ( + f"only the permutation scan may be unbounded, got {unbounded}" + ) diff --git a/python/src/permutation.rs b/python/src/permutation.rs index 4dc49cfd3..24ca2b5a3 100644 --- a/python/src/permutation.rs +++ b/python/src/permutation.rs @@ -268,7 +268,9 @@ impl PyPermutationReader { .await .infer_error()? } else { - PermutationReader::identity(base_table).await + PermutationReader::identity(base_table) + .await + .infer_error()? }; Ok(Self::from_reader(reader)) }) diff --git a/rust/lancedb/src/dataloader/permutation/builder.rs b/rust/lancedb/src/dataloader/permutation/builder.rs index 6b4ae2303..a3700d0ef 100644 --- a/rust/lancedb/src/dataloader/permutation/builder.rs +++ b/rust/lancedb/src/dataloader/permutation/builder.rs @@ -160,9 +160,10 @@ impl PermutationBuilder { self } - async fn sort_by_split_id( + async fn sort_by_column( &self, data: SendableRecordBatchStream, + column: &str, ) -> Result { let memory_limit = std::env::var("LANCEDB_PERM_BUILDER_MEMORY_LIMIT") .unwrap_or_else(|_| DEFAULT_MEMORY_LIMIT.to_string()) @@ -188,25 +189,26 @@ impl PermutationBuilder { let df = ctx .read_one_shot(data.into_df_stream()) .map_err(|e| Error::Other { - message: format!("Failed to setup sort by split id: {}", e), + message: format!("Failed to setup sort by {}: {}", column, e), source: Some(e.into()), })?; let df_stream = df - .sort_by(vec![col(SPLIT_ID_COLUMN)]) + .sort_by(vec![col(column)]) .map_err(|e| Error::Other { - message: format!("Failed to plan sort by split id: {}", e), + message: format!("Failed to plan sort by {}: {}", column, e), source: Some(e.into()), })? .execute_stream() .await .map_err(|e| Error::Other { - message: format!("Failed to sort by split id: {}", e), + message: format!("Failed to sort by {}: {}", column, e), source: Some(e.into()), })?; + let column = column.to_string(); let schema = df_stream.schema(); - let stream = df_stream.map_err(|e| Error::Other { - message: format!("Failed to execute sort by split id: {}", e), + let stream = df_stream.map_err(move |e| Error::Other { + message: format!("Failed to execute sort by {}: {}", column, e), source: Some(e.into()), }); Ok(Box::pin(SimpleRecordBatchStream { schema, stream })) @@ -238,7 +240,25 @@ impl PermutationBuilder { /// Builds the permutation table and stores it in the given database. pub async fn build(self) -> Result { - // First pass, apply filter and load row ids + // Unflushed rows have no row id, so a permutation cannot address them. + match self.base_table.base_table().get_lsm_write_spec().await { + Ok(Some(_)) => { + return Err(Error::NotSupported { + message: "the data loader does not support tables with an LSM write \ + spec: rows that have not been flushed to the base table \ + have no row id, so a permutation cannot reference them" + .to_string(), + }); + } + Ok(None) => {} + // No LSM write path means no spec. + Err(Error::NotSupported { .. }) => {} + Err(err) => return Err(err), + } + + // 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. + // TODO: pin the version resolved here; remote does not implement Lazy. let mut rows = self.base_table.query().select(Select::columns(&[ROW_ID])); if let Some(filter) = &self.config.filter { @@ -263,6 +283,12 @@ impl PermutationBuilder { // Apply splits let rows = rows.execute().await?; + // Splits are assigned by position, so the scan has to arrive in a fixed order. + let rows = if self.base_table.base_table().scan_order_is_deterministic() { + rows + } else { + self.sort_by_column(rows, ROW_ID).await? + }; let split_data = splitter.apply(rows, num_rows).await?; // Shuffle data if requested @@ -284,7 +310,7 @@ impl PermutationBuilder { needs_sort |= !matches!(self.config.shuffle_strategy, ShuffleStrategy::None); let sorted = if needs_sort { - self.sort_by_split_id(shuffled).await? + self.sort_by_column(shuffled, SPLIT_ID_COLUMN).await? } else { shuffled }; @@ -367,6 +393,22 @@ mod tests { ); } + #[tokio::test] + async fn test_native_scan_order_is_deterministic() { + 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("col_a", lance_datagen::array::step::()) + .into_ldb_stream(RowCount::from(10), BatchCount::from(1)); + let table = db.create_table("t", data).execute().await.unwrap(); + + // Native tables skip the canonicalizing sort; remote does not. + assert!(table.base_table().scan_order_is_deterministic()); + } + #[tokio::test] async fn test_permutation_builder() { let temp_dir = tempfile::tempdir().unwrap(); @@ -416,4 +458,48 @@ mod tests { 283 ); } + + /// Rows that have not been flushed to the base table have no row id, so a + /// permutation cannot reference them. Reading the base table alone would drop + /// them from training without saying so, so the table is refused instead. + #[tokio::test] + async fn test_permutation_rejects_lsm_write_spec() { + use crate::table::LsmWriteSpec; + use arrow_array::{Int32Array, RecordBatchIterator}; + use arrow_schema::{DataType, Field, Schema}; + + // MemWAL needs a real dataset directory and a non-nullable primary key. + let temp_dir = tempfile::tempdir().unwrap(); + let db = connect(temp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("idx", DataType::Int32, false)])); + let batch = arrow_array::RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![0, 1, 2, 3]))], + ) + .unwrap(); + let reader: Box = + Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone())); + let table = db.create_table("tbl", reader).execute().await.unwrap(); + + // Without a spec the build succeeds. + PermutationBuilder::new(table.clone()) + .build() + .await + .unwrap(); + + table.set_unenforced_primary_key(["idx"]).await.unwrap(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + + let err = PermutationBuilder::new(table).build().await.unwrap_err(); + assert!( + err.to_string().contains("LSM write spec"), + "expected the pre-check to refuse the table, got: {err}" + ); + } } diff --git a/rust/lancedb/src/dataloader/permutation/reader.rs b/rust/lancedb/src/dataloader/permutation/reader.rs index afe79b0ad..6da92e986 100644 --- a/rust/lancedb/src/dataloader/permutation/reader.rs +++ b/rust/lancedb/src/dataloader/permutation/reader.rs @@ -97,8 +97,10 @@ impl PermutationReader { Self::inner_new(base_table, Some(permutation_table), split).await } - pub async fn identity(base_table: Arc) -> Self { - Self::inner_new(base_table, None, 0).await.unwrap() + /// A reader over the base table in storage order, with no permutation. + /// Fallible because construction counts the base table. + pub async fn identity(base_table: Arc) -> Result { + Self::inner_new(base_table, None, 0).await } /// Validates the limit and offset and returns the number of rows that will be read @@ -487,7 +489,13 @@ impl PermutationReader { pub async fn output_schema(&self, selection: Select) -> Result { let table = Table::from(self.base_table.clone()); - table.query().select(selection).output_schema().await + // limit(1) because some table types execute the query to get its schema + table + .query() + .select(selection) + .limit(1) + .output_schema() + .await } pub fn count_rows(&self) -> u64 { @@ -779,7 +787,9 @@ mod tests { .into_mem_table("tbl", RowCount::from(10), BatchCount::from(1)) .await; - let reader = PermutationReader::identity(base_table.base_table().clone()).await; + let reader = PermutationReader::identity(base_table.base_table().clone()) + .await + .unwrap(); // With no permutation table, take_offsets uses the base table directly let offsets = vec![0, 2, 4, 6]; @@ -961,7 +971,9 @@ mod tests { .into_mem_table("tbl", RowCount::from(10), BatchCount::from(1)) .await; - let reader = PermutationReader::identity(base_table.base_table().clone()).await; + let reader = PermutationReader::identity(base_table.base_table().clone()) + .await + .unwrap(); let batch = reader.take_offsets(&[], Select::All).await.unwrap(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 14b869d90..f98d4c8dd 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -4303,6 +4303,19 @@ mod tests { write_ipc_stream_uncompressed(&one_row_blob_batch(column)) } + #[tokio::test] + async fn test_remote_scan_order_is_not_deterministic() { + // A distributed scan answers in no fixed order, so callers that assign meaning + // to row position have to sort for themselves. + let table = Table::new_with_handler("my_table", |_| { + http::Response::builder() + .status(200) + .body(Vec::new()) + .unwrap() + }); + assert!(!table.base_table().scan_order_is_deterministic()); + } + #[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.rs b/rust/lancedb/src/table.rs index e1dd942db..096d60345 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -789,6 +789,13 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { async fn checkout_tag(&self, tag: &str) -> Result<()>; /// Checkout the latest version of the table. async fn checkout_latest(&self) -> Result<()>; + /// Whether repeated identical scans return rows in the same order. + /// + /// Callers that assign meaning to a row's position must order the results + /// themselves when this is false. Defaults to false so a table type opts in. + fn scan_order_is_deterministic(&self) -> bool { + false + } /// Restore the table to the currently checked out version. async fn restore(&self) -> Result<()>; /// List the versions of the table. @@ -2996,6 +3003,12 @@ impl BaseTable for NativeTable { self } + /// Lance scans fragments in order (`Scanner::ordered` defaults to true, and we + /// never clear it), so repeated identical scans agree. + fn scan_order_is_deterministic(&self) -> bool { + true + } + fn name(&self) -> &str { self.name.as_str() }