bug(python): exact KNN results remain globally unsorted after #7868 due to an internal CoalescePartitionsExec
LanceDB version
0.40.0b2
What happened?
This is the behavior same as issue ID 3669 in lancedb: LanceDB issue #3669 .
The query returns 9,333 rows. Each output batch is internally sorted, but occasionally the batches are concatenated in the wrong order, 1141 farther rows + 8192 nearer rows. It is expected that the nearer rows are arranged first, followed by the farther rows. This problem is not deterministic, and it may require multiple runs of the triggering testcase to trigger the bug.
When run with LANCE_CPU_COUNT=1, the error is not reproduced. When run with LANCE_CPU_COUNT larger than 1, especially around 2-10, the error is more likely to be produce. Below are the query plans for one cpu and 10 cpu
Query plan for LANCE_CPU_COUNT=1
ProjectionExec: expr=[f0@4 as f0, f1@5 as f1, f2@0 as f2, f3@6 as f3, f4@7 as f4, f5@8 as f5, f6@9 as f6, f7@1 as f7, _distance@3 as _distance]
LanceRead: uri=tmp/lancedb_trigger_15__t5xjs7e/test_collection.lance/data, projection=[f0, f1, f3, f4, f5, f6], source=stream(_rowid)
GlobalLimitExec: skip=0, fetch=9336
FilterExec: _distance@3 IS NOT NULL
SortExec: TopK(fetch=9336), expr=[_distance@3 ASC NULLS LAST, _rowid@2 ASC NULLS LAST], preserve_partitioning=[false]
KNNVectorDistance: metric=l2
FilterExec: f2@0 NOT LIKE %At%@v%dN%6u
LanceRead: uri=tmp/lancedb_trigger_15__t5xjs7e/test_collection.lance/data, projection=[f2, f7], num_fragments=2, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=f2 NOT LIKE Utf8("%At%@v%dN%6u"), refine_filter=f2 NOT LIKE Utf8("%At%@v%dN%6u")
Query plan for LANCE_CPU_COUNT=10
ProjectionExec: expr=[f0@4 as f0, f1@5 as f1, f2@0 as f2, f3@6 as f3, f4@7 as f4, f5@8 as f5, f6@9 as f6, f7@1 as f7, _distance@3 as _distance]
LanceRead: uri=tmp/lancedb_trigger_15_uzj6fdkn/test_collection.lance/data, projection=[f0, f1, f3, f4, f5, f6], source=stream(_rowid)
GlobalLimitExec: skip=0, fetch=9336
CoalescePartitionsExec
FilterExec: _distance@3 IS NOT NULL
RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
SortExec: TopK(fetch=9336), expr=[_distance@3 ASC NULLS LAST, _rowid@2 ASC NULLS LAST], preserve_partitioning=[false]
KNNVectorDistance: metric=l2
CoalescePartitionsExec
FilterExec: f2@0 NOT LIKE %At%@v%dN%6u
RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
LanceRead: uri=tmp/lancedb_trigger_15_uzj6fdkn/test_collection.lance/data, projection=[f2, f7], num_fragments=2, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=f2 NOT LIKE Utf8("%At%@v%dN%6u"), refine_filter=f2 NOT LIKE Utf8("%At%@v%dN%6u")
The fix for #3669 was on the outer most level, for the case where partition is larger than 1, but the CoalescePartitionsExec step concatenates the blocks into one partition out of order, bypassing the fix in #3669.
Are there known steps to reproduce?
Run the same trigger as #3669, with LANCE_CPU_COUNT=1 and =10 to observe the bug.
Source: lancedb/lancedb