diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index d611fe8a5..7cb622d46 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -17,11 +17,15 @@ use arrow::array::{AsArray, FixedSizeListBuilder, Float32Builder}; use arrow::datatypes::{Float32Type, UInt8Type}; use arrow_array::Array; use arrow_schema::{DataType, Schema}; +use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{Column, DataFusionError, SchemaError}; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion_physical_plan::limit::GlobalLimitExec; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::UnionExec; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use lance::dataset::mem_wal::DatasetMemWalExt; use lance::dataset::scanner::DatasetRecordBatchStream; use lance::dataset::scanner::Scanner; @@ -382,10 +386,42 @@ pub async fn create_plan( scanner.order_by(Some(order_by.clone()))?; } - scanner + let plan = scanner .create_plan() .await - .map_err(|error| enrich_lance_field_not_found(error, schema)) + .map_err(|error| enrich_lance_field_not_found(error, schema))?; + preserve_ordered_global_limits(plan) +} + +/// Preserve an ordered input when Lance's distribution optimizer coalesces it for a limit. +/// +/// Exact KNN plans sort row ids before applying their global limit. At parallelism greater +/// than one, `EnforceDistribution` can insert a [`CoalescePartitionsExec`] immediately below +/// that limit. A plain coalesce concatenates its ordered input partitions in completion order, +/// so a later materialization step can return individually sorted batches in the wrong global +/// order. Replace only this limit-bound coalesce with an order-preserving merge. +fn preserve_ordered_global_limits(plan: Arc) -> Result> { + Ok(plan + .transform_up(|plan| { + let Some(limit) = plan.downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + let Some(coalesce) = limit.input().downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + let input = coalesce.input(); + let Some(ordering) = input.output_ordering().cloned() else { + return Ok(Transformed::no(plan)); + }; + + let merge: Arc = Arc::new( + SortPreservingMergeExec::new(ordering, input.clone()).with_fetch(coalesce.fetch()), + ); + let mut replacement = GlobalLimitExec::new(merge, limit.skip(), limit.fetch()); + replacement.set_required_ordering(limit.required_ordering().clone()); + Ok(Transformed::yes(Arc::new(replacement))) + })? + .data) } /// Replace DataFusion's top-level field candidates with qualified leaf paths. @@ -1024,6 +1060,108 @@ mod tests { assert_eq!(count, 2); // 4 and 5 } + #[tokio::test(flavor = "multi_thread")] + async fn test_exact_knn_large_limit_preserves_global_order() { + use crate::connect; + use arrow_array::cast::AsArray; + use arrow_schema::{DataType, Field, Schema}; + use lance_index::vector::DIST_COL; + + const ROWS: usize = 20_000; + const LIMIT: usize = 12_000; + + fn has_order_losing_limit(plan: &Arc) -> bool { + if let Some(limit) = plan.downcast_ref::() + && let Some(coalesce) = limit.input().downcast_ref::() + && coalesce.input().output_ordering().is_some() + { + return true; + } + plan.children() + .iter() + .any(|child| has_order_losing_limit(child)) + } + + fn has_order_preserving_limit(plan: &Arc) -> bool { + if let Some(limit) = plan.downcast_ref::() + && limit + .input() + .downcast_ref::() + .is_some() + { + return true; + } + plan.children() + .iter() + .any(|child| has_order_preserving_limit(child)) + } + + let conn = connect("memory://").execute().await.unwrap(); + let ids = Arc::new(Int32Array::from_iter_values(0..ROWS as i32)); + let vectors = Arc::new(fixed_size_list_array( + (0..ROWS).flat_map(|value| [value as f32, 0.0]).collect(), + 2, + )); + // This column is materialized after the KNN limit, matching the plan shape + // that exposed the internal coalesce in the reported Python query. + let payload = Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + "payload", ROWS, + ))); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("vector", vectors.data_type().clone(), false), + Field::new("payload", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new(schema, vec![ids, vectors, payload]).unwrap(); + let table = conn + .create_table("ordered_exact_knn", batch) + .execute() + .await + .unwrap(); + let dataset = table.as_native().unwrap().dataset.get().await.unwrap(); + + let mut scanner = dataset.scan(); + let query = Float32Array::from(vec![0.0, 0.0]); + scanner.nearest("vector", &query, LIMIT).unwrap(); + scanner.limit(Some(LIMIT as i64), None).unwrap(); + scanner.target_parallelism(8); + let plan = scanner.create_plan().await.unwrap(); + let displayed = + datafusion_physical_plan::display::DisplayableExecutionPlan::new(plan.as_ref()) + .indent(true) + .to_string(); + + assert!( + has_order_losing_limit(&plan), + "test setup must contain an ordered input coalesced below a global limit:\n{displayed}" + ); + let plan = preserve_ordered_global_limits(plan).unwrap(); + assert!(!has_order_losing_limit(&plan)); + assert!(has_order_preserving_limit(&plan)); + + let stream = execute_plan(plan, Default::default()).unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + let mut distances = Vec::with_capacity(LIMIT); + for batch in batches { + distances.extend_from_slice( + batch + .column_by_name(DIST_COL) + .unwrap() + .as_primitive::() + .values(), + ); + } + assert_eq!(distances.len(), LIMIT); + for pair in distances.windows(2) { + assert!( + pair[0] <= pair[1], + "exact KNN distances must be globally ordered, found {} before {}", + pair[0], + pair[1] + ); + } + } + #[tokio::test] async fn test_missing_filter_field_lists_nested_fields_in_local_planners() { use crate::connect;