mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-14 15:52:26 +00:00
fix: preserve duplicate take offsets (#4024)
## Summary - preserve repeated table offsets without adding a public ordering guarantee - retain exact requested ordering in identity and persisted permutations - cover local, projected, multi-batch, and mocked-remote query paths ## Root cause Take queries lowered offsets to a set-like IN predicate and discarded repeated occurrences. Persisted permutation loading also compared the distinct base-table result count with the requested occurrence count, rejecting repeated row IDs before its existing reordering step could expand them. ## Fix The shared take-query path now deduplicates the predicate for efficient lookup, requests row-offset metadata internally, and expands each matching row to the requested multiplicity in backend result order. An internal opt-in keeps exact requested order for identity PermutationReader reads, while persisted permutations continue using their existing ordering map. ## Validation - cargo test --quiet --features remote --tests - cargo check --quiet --features remote --tests --examples - cargo clippy --quiet --features remote --tests --examples - targeted Python local and mocked-remote regression tests - exact issue reproduction Fixes #2820 <!-- lance-gatekeeper-fix:v1 agent=75acf840afa6f4be4bff98b567b504bd generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
co-authored by
Xuanwo
parent
5cbd979455
commit
19232f9c50
@@ -1923,6 +1923,21 @@ def test_take_queries(tmp_path):
|
||||
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,
|
||||
]
|
||||
|
||||
# Converting a take builder to its serializable query representation must
|
||||
# retain occurrence metadata and execute with the same multiplicity.
|
||||
query = table.take_offsets([5, 2, 5, 17]).select(["idx"]).to_query_object()
|
||||
assert query.take_offsets == [5, 2, 5, 17]
|
||||
converted = table._execute_query(query).read_all()
|
||||
assert sorted(converted["idx"].to_pylist()) == [2, 5, 5, 17]
|
||||
|
||||
# Take by row id
|
||||
assert list(
|
||||
sorted(table.take_row_ids([5, 2, 17]).to_pandas()["idx"].to_list())
|
||||
|
||||
@@ -479,24 +479,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": 0},
|
||||
{"a": 2},
|
||||
{"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():
|
||||
|
||||
Reference in New Issue
Block a user