diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 913ab5289..284793c5d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1489,9 +1489,9 @@ class Table(ABC): Offsets are mostly useful for sampling as the set of all valid offsets is easily known in advance to be [0, len(table)). - No guarantees are made regarding the order in which results are returned. If - you desire an output order that matches the order of the given offsets, you will - need to add the row offset column to the output and align it yourself. + Results are returned in the same order as the given offsets. Repeated offsets + produce repeated rows, which makes this method suitable for sampling with + replacement. Parameters ---------- @@ -6291,6 +6291,9 @@ class AsyncTable: Offsets are mostly useful for sampling as the set of all valid offsets is easily known in advance to be [0, len(table)). + Results are returned in the same order as the given offsets, including repeated + occurrences. + Parameters ---------- offsets: list[int] diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index d2629d1a8..a6b7e1621 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -1891,6 +1891,14 @@ def test_take_queries(tmp_path): 17, ] + # Duplicate offsets are occurrences, not set members, and preserve input order. + assert table.take_offsets([5, 2, 5, 17]).to_pandas()["idx"].to_list() == [ + 5, + 2, + 5, + 17, + ] + # Take by row id assert list( sorted(table.take_row_ids([5, 2, 17]).to_pandas()["idx"].to_list()) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 13ffc4415..6ff2dfa62 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -480,24 +480,49 @@ def test_remote_permutation_is_picklable(): match = re.search( r"_rowoffset\s+in\s+\((.*?)\)", body["filter"], re.IGNORECASE ) - offsets = [int(o.strip()) for o in match.group(1).split(",")] + offsets = list( + dict.fromkeys(int(o.strip()) for o in match.group(1).split(",")) + ) else: offsets = list(range(len(rows))) - table = pa.table({"a": [rows[offset] for offset in offsets]}) + columns = body.get("columns") or ["a"] + table = pa.table( + { + column: ( + [rows[offset] for offset in offsets] + if column == "a" + else offsets + ) + for column in columns + } + ) request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.file") request.end_headers() with pa.ipc.new_file(request.wfile, schema=table.schema) as writer: - writer.write_table(table) + writer.write_table(table, max_chunksize=2) else: request.send_response(404) request.end_headers() with mock_lancedb_connection(handler) as db: - permutation = Permutation.identity(db.open_table("test")) + table = db.open_table("test") + assert table.take_offsets([0, 2, 0, 4]).to_list() == [ + {"a": 0}, + {"a": 2}, + {"a": 0}, + {"a": 4}, + ] + + permutation = Permutation.identity(table) restored = pickle.loads(pickle.dumps(permutation)) - assert restored.__getitems__([0, 2, 4]) == [{"a": 0}, {"a": 2}, {"a": 4}] + assert restored.__getitems__([0, 2, 0, 4]) == [ + {"a": 0}, + {"a": 2}, + {"a": 0}, + {"a": 4}, + ] def test_create_table_exist_ok(): diff --git a/rust/lancedb/src/dataloader/permutation/reader.rs b/rust/lancedb/src/dataloader/permutation/reader.rs index 6da92e986..93c7b1120 100644 --- a/rust/lancedb/src/dataloader/permutation/reader.rs +++ b/rust/lancedb/src/dataloader/permutation/reader.rs @@ -28,7 +28,7 @@ use lance::io::RecordBatchStream; use lance_arrow::RecordBatchExt; use lance_core::ROW_ID; use lance_core::error::LanceOptionExt; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; /// Reads a permutation of a source table based on row IDs stored in a separate table @@ -199,7 +199,14 @@ impl PermutationReader { .expect_ok()? .values(); - let in_list: Vec = row_ids.iter().map(|id| lit(*id)).collect(); + let mut unique_row_ids = HashSet::with_capacity(num_rows); + let in_list: Vec = row_ids + .iter() + .copied() + .filter(|row_id| unique_row_ids.insert(*row_id)) + .map(lit) + .collect(); + let num_unique_row_ids = unique_row_ids.len(); let base_query = QueryRequest { filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))), @@ -212,7 +219,7 @@ impl PermutationReader { .query( &AnyQuery::Query(base_query), QueryExecutionOptions { - max_batch_length: num_rows as u32, + max_batch_length: num_unique_row_ids as u32, ..Default::default() }, ) @@ -227,9 +234,9 @@ impl PermutationReader { }); } - if batches.iter().map(|b| b.num_rows()).sum::() != num_rows { + if batches.iter().map(|b| b.num_rows()).sum::() != num_unique_row_ids { return Err(Error::InvalidInput { - message: "Base table returned different number of rows than the number of row IDs" + message: "Base table returned a different number of rows than the number of unique row IDs" .to_string(), }); } @@ -712,10 +719,10 @@ mod tests { .unwrap(); // Take offsets in reverse order and verify returned rows match that order - let offsets = vec![5, 3, 1, 0]; + let offsets = vec![5, 3, 5, 1, 0]; let batch = reader.take_offsets(&offsets, Select::All).await.unwrap(); - assert_eq!(batch.num_rows(), 4); + assert_eq!(batch.num_rows(), 5); let idx_values = batch .column(0) @@ -729,6 +736,52 @@ mod tests { assert_eq!(idx_values, expected); } + #[tokio::test] + async fn test_take_offsets_preserves_repeated_rows_in_permutation() { + let base_table = lance_datagen::gen_batch() + .col("idx", lance_datagen::array::step::()) + .into_mem_table("tbl", RowCount::from(5), BatchCount::from(1)) + .await; + let base_row_ids = collect_column::(&base_table, "_rowid").await; + let permutation_row_ids = vec![ + base_row_ids[3], + base_row_ids[1], + base_row_ids[3], + base_row_ids[2], + ]; + let permutation_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("row_id", DataType::UInt64, false), + Field::new(SPLIT_ID_COLUMN, DataType::UInt64, false), + ])), + vec![ + Arc::new(UInt64Array::from(permutation_row_ids)), + Arc::new(UInt64Array::from(vec![0; 4])), + ], + ) + .unwrap(); + let permutation_table = virtual_table("row_ids", &permutation_batch).await; + let reader = PermutationReader::try_from_tables( + base_table.base_table().clone(), + permutation_table.base_table().clone(), + 0, + ) + .await + .unwrap(); + + let batch = reader + .take_offsets(&[0, 1, 2, 3], Select::All) + .await + .unwrap(); + let idx_values = batch + .column(0) + .as_primitive::() + .values() + .to_vec(); + + assert_eq!(idx_values, vec![3, 1, 3, 2]); + } + #[tokio::test] async fn test_take_offsets_with_column_selection() { let (base_table, row_ids_table, row_ids) = setup_permutation_tables(10).await; @@ -792,17 +845,17 @@ mod tests { .unwrap(); // With no permutation table, take_offsets uses the base table directly - let offsets = vec![0, 2, 4, 6]; + let offsets = vec![0, 2, 0, 4, 6]; let batch = reader.take_offsets(&offsets, Select::All).await.unwrap(); - assert_eq!(batch.num_rows(), 4); + assert_eq!(batch.num_rows(), 5); let idx_values = batch .column(0) .as_primitive::() .values() .to_vec(); - assert_eq!(idx_values, vec![0, 2, 4, 6]); + assert_eq!(idx_values, vec![0, 2, 0, 4, 6]); } #[tokio::test] diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index b2c5fefbe..82d91a02d 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1,11 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::{future::Future, time::Duration}; use arrow::compute::concat_batches; -use arrow_array::{Array, Float16Array, Float32Array, Float64Array, RecordBatch, make_array}; +use arrow_array::{ + Array, Float16Array, Float32Array, Float64Array, RecordBatch, UInt64Array, + cast::AsArray, + make_array, + types::{Int64Type, UInt64Type}, +}; use arrow_schema::{DataType, SchemaRef}; use datafusion_expr::{Expr, col, lit}; use datafusion_physical_plan::ExecutionPlan; @@ -14,6 +20,7 @@ use half::f16; /// Re-export Lance ColumnOrdering type for use in query ordering pub use lance::dataset::scanner::ColumnOrdering; use lance::dataset::{ROW_ID, scanner::DatasetRecordBatchStream}; +use lance::io::RecordBatchStream; use lance_arrow::RecordBatchExt; use lance_datafusion::exec::execute_plan; use lance_index::scalar::FullTextSearchQuery; @@ -1531,6 +1538,7 @@ impl HasQuery for VectorQuery { pub struct TakeQuery { parent: Arc, request: QueryRequest, + offsets: Option>, } impl TakeQuery { @@ -1538,7 +1546,13 @@ impl TakeQuery { /// /// See [`crate::Table::take_offsets`] for more details. pub fn from_offsets(parent: Arc, offsets: Vec) -> Self { - let in_list: Vec = offsets.iter().map(|o| lit(*o)).collect(); + let mut seen = HashSet::with_capacity(offsets.len()); + let in_list: Vec = offsets + .iter() + .copied() + .filter(|offset| seen.insert(*offset)) + .map(lit) + .collect(); Self { parent, request: QueryRequest { @@ -1547,6 +1561,7 @@ impl TakeQuery { )), ..Default::default() }, + offsets: Some(offsets), } } @@ -1561,9 +1576,145 @@ impl TakeQuery { filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))), ..Default::default() }, + offsets: None, } } + async fn request_with_row_offset(&self) -> Result<(QueryRequest, String, bool)> { + const ROW_OFFSET: &str = "_rowoffset"; + const INTERNAL_ROW_OFFSET: &str = "__lancedb_take_row_offset"; + + let mut request = self.request.clone(); + let (ordering_column, drop_ordering_column) = match &mut request.select { + Select::All => { + let mut columns = self + .parent + .schema() + .await? + .fields() + .iter() + .map(|field| field.name().clone()) + .collect::>(); + columns.push(ROW_OFFSET.to_string()); + request.select = Select::Columns(columns); + (ROW_OFFSET.to_string(), true) + } + Select::Columns(columns) => { + if columns.iter().any(|column| column == ROW_OFFSET) { + (ROW_OFFSET.to_string(), false) + } else { + columns.push(ROW_OFFSET.to_string()); + (ROW_OFFSET.to_string(), true) + } + } + Select::Dynamic(columns) => { + let mut ordering_column = INTERNAL_ROW_OFFSET.to_string(); + while columns.iter().any(|(name, _)| name == &ordering_column) { + ordering_column.push('_'); + } + columns.push((ordering_column.clone(), ROW_OFFSET.to_string())); + (ordering_column, true) + } + Select::Expr(columns) => { + let mut ordering_column = INTERNAL_ROW_OFFSET.to_string(); + while columns.iter().any(|(name, _)| name == &ordering_column) { + ordering_column.push('_'); + } + columns.push((ordering_column.clone(), col(ROW_OFFSET))); + (ordering_column, true) + } + }; + + Ok((request, ordering_column, drop_ordering_column)) + } + + async fn execute_offsets( + &self, + offsets: &[u64], + options: QueryExecutionOptions, + ) -> Result { + let max_batch_length = options.max_batch_length as usize; + let (request, ordering_column, drop_ordering_column) = + self.request_with_row_offset().await?; + let query = AnyQuery::Query(request); + let data = self + .parent + .clone() + .query(&query, options.without_output_batch_length_limit()) + .await?; + let schema = data.schema(); + let batches = data.try_collect::>().await?; + let batch = if batches.is_empty() { + RecordBatch::new_empty(schema) + } else { + concat_batches(&schema, &batches)? + }; + + let actual_offsets = + batch + .column_by_name(&ordering_column) + .ok_or_else(|| Error::Schema { + message: format!( + "take query result did not include ordering column '{ordering_column}'" + ), + })?; + let actual_offsets = match actual_offsets.data_type() { + DataType::UInt64 => actual_offsets + .as_primitive::() + .values() + .to_vec(), + DataType::Int64 => actual_offsets + .as_primitive::() + .values() + .iter() + .map(|offset| { + u64::try_from(*offset).map_err(|_| Error::Schema { + message: format!( + "take query ordering column '{ordering_column}' contained a negative offset" + ), + }) + }) + .collect::>>()?, + data_type => { + return Err(Error::Schema { + message: format!( + "take query ordering column '{ordering_column}' had unsupported type {data_type}" + ), + }); + } + }; + + let ordering = actual_offsets + .iter() + .copied() + .enumerate() + .map(|(index, offset)| (offset, index as u64)) + .collect::>(); + // Missing offsets retain the filter-based behavior of returning no row. Every + // occurrence of an offset that was found is restored in the requested order. + let desired_order = offsets + .iter() + .filter_map(|offset| ordering.get(offset).copied()) + .collect::>(); + + let mut ordered_batch = if desired_order.len() == batch.num_rows() + && desired_order + .iter() + .enumerate() + .all(|(index, desired)| *desired == index as u64) + { + batch + } else { + arrow_select::take::take_record_batch(&batch, &UInt64Array::from(desired_order))? + }; + + if drop_ordering_column { + ordered_batch = ordered_batch.drop_column(&ordering_column)?; + } + + Ok(single_batch_stream(ordered_batch, max_batch_length)) + } + /// Convert the `TakeQuery` into a `QueryRequest`. pub fn into_request(self) -> QueryRequest { self.request @@ -1624,6 +1775,10 @@ impl ExecutableQuery for TakeQuery { &self, options: QueryExecutionOptions, ) -> Result { + if let Some(offsets) = &self.offsets { + return self.execute_offsets(offsets, options).await; + } + let query = AnyQuery::Query(self.request.clone()); Ok(SendableRecordBatchStream::from( self.parent.clone().query(&query, options).await?, @@ -2657,6 +2812,40 @@ mod tests { assert_eq!(results[0].num_columns(), 1); } + #[tokio::test] + async fn test_take_offsets_preserves_duplicate_order() { + let tmp_dir = tempdir().unwrap(); + let table = make_test_table(&tmp_dir).await; + + let results = table + .take_offsets(vec![5, 1, 5, 17]) + .select(Select::Columns(vec!["id".to_string()])) + .execute_with_options(QueryExecutionOptions { + max_batch_length: 2, + ..Default::default() + }) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(results.len(), 2); + assert!(results.iter().all(|batch| batch.num_columns() == 1)); + let ids = results + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect::>(); + assert_eq!(ids, vec![5, 1, 5, 17]); + } + #[tokio::test] async fn test_take_row_ids() { let tmp_dir = tempdir().unwrap(); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 096d60345..73ebfd619 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1602,9 +1602,9 @@ impl Table { /// Offsets are useful for sampling as the set of all valid offsets is easily /// known in advance to be [0, len(table)). /// - /// No guarantees are made regarding the order in which results are returned. If you - /// desire an output order that matches the order of the given offsets, you will need - /// to add the row offset column to the output and align it yourself. + /// Results are returned in the same order as the given offsets. Repeated offsets + /// produce repeated rows, which makes this method suitable for sampling with + /// replacement. /// /// Parameters /// ----------