diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 284793c5d..053cd26ea 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)). - 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. + No guarantees are made regarding the order in which results are returned. + Repeated offsets produce repeated rows, which makes this method suitable for + sampling with replacement. Parameters ---------- @@ -6291,8 +6291,8 @@ 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. + No guarantees are made regarding the order in which results are returned. + Repeated offsets produce repeated rows. Parameters ---------- diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index a6b7e1621..1a709a52e 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -1891,13 +1891,10 @@ 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, - ] + # Duplicate offsets are occurrences, not set members. Ordering is unspecified. + assert sorted( + table.take_offsets([5, 2, 5, 17]).to_pandas()["idx"].to_list() + ) == [2, 5, 5, 17] # Take by row id assert list( diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 6ff2dfa62..00e6a3983 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -510,8 +510,8 @@ def test_remote_permutation_is_picklable(): table = db.open_table("test") assert table.take_offsets([0, 2, 0, 4]).to_list() == [ {"a": 0}, - {"a": 2}, {"a": 0}, + {"a": 2}, {"a": 4}, ] diff --git a/rust/lancedb/src/dataloader/permutation/reader.rs b/rust/lancedb/src/dataloader/permutation/reader.rs index 93c7b1120..03dacb9c4 100644 --- a/rust/lancedb/src/dataloader/permutation/reader.rs +++ b/rust/lancedb/src/dataloader/permutation/reader.rs @@ -476,6 +476,7 @@ impl PermutationReader { let table = Table::from(self.base_table.clone()); let batches = table .take_offsets(offsets.to_vec()) + .preserve_order() .select(selection.clone()) .execute() .await? diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 6f588d2c5..9afb35356 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1535,6 +1535,7 @@ fn restore_take_batch( offsets: &[u64], ordering_column: &str, drop_ordering_column: bool, + preserve_order: bool, ) -> Result { let actual_offsets = batch .column_by_name(ordering_column) @@ -1569,18 +1570,33 @@ fn restore_take_batch( } }; - 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 desired_order = Vec::with_capacity(offsets.len()); + if preserve_order { + 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. + desired_order.extend( + offsets + .iter() + .filter_map(|offset| ordering.get(offset).copied()), + ); + } else { + let mut occurrences = HashMap::with_capacity(offsets.len()); + for offset in offsets { + *occurrences.entry(*offset).or_insert(0) += 1; + } + // Public take queries do not guarantee output order. Preserve the lookup's + // existing order and only restore the multiplicity of each matching row. + for (index, offset) in actual_offsets.iter().enumerate() { + if let Some(count) = occurrences.remove(offset) { + desired_order.extend(std::iter::repeat_n(index as u64, count)); + } + } + } let mut ordered_batch = if desired_order.len() == batch.num_rows() && desired_order @@ -1603,14 +1619,16 @@ fn restore_take_batch( /// Restores the logical offset occurrence sequence above the physical lookup plan. /// /// The lookup plan returns each matching row at most once. This operator collects -/// those rows, expands duplicates, and emits one partition in the caller's offset -/// order. Pagination must remain above this operator so it applies to occurrences. +/// those rows, expands duplicates, and emits one partition. It preserves lookup order +/// unless the caller explicitly requests offset order. Pagination must remain above +/// this operator so it applies to occurrences. #[derive(Debug)] struct TakeRestoreExec { input: Arc, offsets: Vec, ordering_column: String, drop_ordering_column: bool, + preserve_order: bool, schema: SchemaRef, properties: Arc, } @@ -1621,6 +1639,7 @@ impl TakeRestoreExec { offsets: Vec, ordering_column: String, drop_ordering_column: bool, + preserve_order: bool, ) -> Result { let schema = if drop_ordering_column { RecordBatch::new_empty(input.schema()) @@ -1641,6 +1660,7 @@ impl TakeRestoreExec { offsets, ordering_column, drop_ordering_column, + preserve_order, schema, properties, }) @@ -1698,6 +1718,7 @@ impl ExecutionPlan for TakeRestoreExec { self.offsets.clone(), self.ordering_column.clone(), self.drop_ordering_column, + self.preserve_order, ) .map_err(|error| DataFusionError::External(Box::new(error)))?; Ok(Arc::new(plan)) @@ -1720,6 +1741,7 @@ impl ExecutionPlan for TakeRestoreExec { let offsets = self.offsets.clone(); let ordering_column = self.ordering_column.clone(); let drop_ordering_column = self.drop_ordering_column; + let preserve_order = self.preserve_order; let stream = stream::once(async move { let batches = input.try_collect::>().await?; let batch = if batches.is_empty() { @@ -1727,8 +1749,14 @@ impl ExecutionPlan for TakeRestoreExec { } else { concat_batches(&input_schema, &batches)? }; - restore_take_batch(batch, &offsets, &ordering_column, drop_ordering_column) - .map_err(|error| DataFusionError::External(Box::new(error))) + restore_take_batch( + batch, + &offsets, + &ordering_column, + drop_ordering_column, + preserve_order, + ) + .map_err(|error| DataFusionError::External(Box::new(error))) }); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -1759,6 +1787,7 @@ pub struct TakeQuery { parent: Arc, request: QueryRequest, offsets: Option>, + preserve_order: bool, } impl TakeQuery { @@ -1782,6 +1811,7 @@ impl TakeQuery { ..Default::default() }, offsets: Some(offsets), + preserve_order: false, } } @@ -1797,9 +1827,19 @@ impl TakeQuery { ..Default::default() }, offsets: None, + preserve_order: false, } } + /// Preserve the requested offset order when restoring duplicate occurrences. + /// + /// This is reserved for readers whose API explicitly guarantees ordering. + pub(crate) fn preserve_order(mut self) -> Self { + debug_assert!(self.offsets.is_some()); + self.preserve_order = true; + self + } + 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"; @@ -1874,6 +1914,7 @@ impl TakeQuery { drop_ordering_column: bool, output_offset: usize, output_limit: Option, + preserve_order: bool, ) -> Result> { let lookup = Arc::new(CoalescePartitionsExec::new(lookup)); let restored: Arc = Arc::new(TakeRestoreExec::try_new( @@ -1881,6 +1922,7 @@ impl TakeQuery { offsets.to_vec(), ordering_column, drop_ordering_column, + preserve_order, )?); if output_offset > 0 || output_limit.is_some() { @@ -1947,6 +1989,7 @@ impl TakeQuery { drop_ordering_column, output_offset, output_limit, + self.preserve_order, ) } @@ -3098,7 +3141,7 @@ mod tests { } #[tokio::test] - async fn test_take_offsets_preserves_duplicate_order() { + async fn test_take_offsets_preserves_duplicate_multiplicity() { let tmp_dir = tempdir().unwrap(); let table = make_test_table(&tmp_dir).await; @@ -3117,7 +3160,7 @@ mod tests { assert_eq!(results.len(), 2); assert!(results.iter().all(|batch| batch.num_columns() == 1)); - let ids = results + let mut ids = results .iter() .flat_map(|batch| { batch @@ -3128,7 +3171,44 @@ mod tests { .to_vec() }) .collect::>(); - assert_eq!(ids, vec![5, 1, 5, 17]); + ids.sort_unstable(); + assert_eq!(ids, vec![1, 5, 5, 17]); + } + + #[test] + fn test_restore_take_batch_only_reorders_when_requested() { + let batch = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int32Array::from(vec![17, 5, 1])) as Arc, + ), + ( + "_rowoffset", + Arc::new(UInt64Array::from(vec![17, 5, 1])) as Arc, + ), + ]) + .unwrap(); + + let restored = + restore_take_batch(batch.clone(), &[5, 1, 5, 17], "_rowoffset", true, false).unwrap(); + assert_eq!( + restored + .column_by_name("id") + .unwrap() + .as_primitive::() + .values(), + &[17, 5, 5, 1] + ); + + let ordered = restore_take_batch(batch, &[5, 1, 5, 17], "_rowoffset", true, true).unwrap(); + assert_eq!( + ordered + .column_by_name("id") + .unwrap() + .as_primitive::() + .values(), + &[5, 1, 5, 17] + ); } #[tokio::test] @@ -3147,13 +3227,15 @@ mod tests { .await .unwrap(); let limited = concat_batches(&limited[0].schema(), &limited).unwrap(); - assert_eq!( + assert_eq!(limited.num_rows(), 3); + assert!( limited .column_by_name("id") .unwrap() .as_primitive::() - .values(), - &[0, 1, 0] + .values() + .iter() + .all(|id| [0, 1, 2].contains(id)) ); let offset = table @@ -3167,13 +3249,15 @@ mod tests { .await .unwrap(); let offset = concat_batches(&offset[0].schema(), &offset).unwrap(); - assert_eq!( + assert_eq!(offset.num_rows(), 3); + assert!( offset .column_by_name("id") .unwrap() .as_primitive::() - .values(), - &[1, 5, 17] + .values() + .iter() + .all(|id| [1, 5, 17].contains(id)) ); } @@ -3197,14 +3281,14 @@ mod tests { .await .unwrap(); let planned = concat_batches(&planned[0].schema(), &planned).unwrap(); - assert_eq!( - planned - .column_by_name("id") - .unwrap() - .as_primitive::() - .values(), - &[5, 1, 5, 17] - ); + let mut ids = planned + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec(); + ids.sort_unstable(); + assert_eq!(ids, vec![1, 5, 5, 17]); } #[tokio::test] diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 3c02ff361..ca59ba30f 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1615,9 +1615,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)). /// - /// 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. + /// No guarantees are made regarding the order in which results are returned. + /// Repeated offsets produce repeated rows, which makes this method suitable for + /// sampling with replacement. /// /// Parameters /// ----------