fix: preserve duplicate take offsets

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]

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())

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():

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<Expr> = row_ids.iter().map(|id| lit(*id)).collect();
+        let mut unique_row_ids = HashSet::with_capacity(num_rows);
+        let in_list: Vec<Expr> = 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::<usize>() != num_rows {
+        if batches.iter().map(|b| b.num_rows()).sum::<usize>() != 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::<Int32Type>())
+            .into_mem_table("tbl", RowCount::from(5), BatchCount::from(1))
+            .await;
+        let base_row_ids = collect_column::<UInt64Type>(&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::<Int32Type>()
+            .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::<Int32Type>()
             .values()
             .to_vec();
-        assert_eq!(idx_values, vec![0, 2, 4, 6]);
+        assert_eq!(idx_values, vec![0, 2, 0, 4, 6]);
     }
 
     #[tokio::test]

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<dyn BaseTable>,
     request: QueryRequest,
+    offsets: Option<Vec<u64>>,
 }
 
 impl TakeQuery {
@@ -1538,7 +1546,13 @@ impl TakeQuery {
     ///
     /// See [`crate::Table::take_offsets`] for more details.
     pub fn from_offsets(parent: Arc<dyn BaseTable>, offsets: Vec<u64>) -> Self {
-        let in_list: Vec<Expr> = offsets.iter().map(|o| lit(*o)).collect();
+        let mut seen = HashSet::with_capacity(offsets.len());
+        let in_list: Vec<Expr> = 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::<Vec<_>>();
+                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<SendableRecordBatchStream> {
+        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::<Vec<_>>().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::<UInt64Type>()
+                .values()
+                .to_vec(),
+            DataType::Int64 => actual_offsets
+                .as_primitive::<Int64Type>()
+                .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::<Result<Vec<_>>>()?,
+            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::<HashMap<_, _>>();
+        // 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::<Vec<_>>();
+
+        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<SendableRecordBatchStream> {
+        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::<Vec<_>>()
+            .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::<Int32Type>()
+                    .values()
+                    .to_vec()
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(ids, vec![5, 1, 5, 17]);
+    }
+
     #[tokio::test]
     async fn test_take_row_ids() {
         let tmp_dir = tempdir().unwrap();

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
     /// ----------