From 03bd909dab8a81ba01d5acad2917d44a824594bd Mon Sep 17 00:00:00 2001 From: Lei Xu Date: Wed, 28 Oct 2026 21:19:44 -0700 Subject: [PATCH] test(python): resolve aliased projections in the torch remote mock server Map dynamic projections by output name to their source columns in the remote torch test mock so permutation take queries preserve the aliased ordering column. Run the regression in the embeddings-installed Python test job and require torch to import so a missing dependency cannot skip coverage. --- .github/workflows/python.yml | 4 ++++ python/python/tests/test_torch.py | 24 +++++++++++++++--------- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index a18ae5ad9..fa6ac4741 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -122,6 +122,10 @@ jobs: run: | pip install --extra-index-url https://pypi.fury.io/lance-format/ --extra-index-url https://pypi.fury.io/lancedb/ -e .[tests,dev,embeddings] pip install mlx + - name: Test remote torch data loader + run: | + python -c "import torch" + pytest python/tests/test_torch.py::test_remote_permutation_dataloader_multiprocessing -q -rs - name: Doctest run: pytest --doctest-modules python/lancedb linux: diff --git a/python/python/tests/test_torch.py b/python/python/tests/test_torch.py index 500378fef..befa3c10f 100644 --- a/python/python/tests/test_torch.py +++ b/python/python/tests/test_torch.py @@ -74,17 +74,23 @@ def _remote_dataset_handler(request): body = json.loads(request.rfile.read(content_len)) offsets = _offsets_from_filter(body.get("filter")) requested_columns = body.get("columns") or ["a"] + # A dynamic projection is sent as {output name: source expression}; take + # queries use it to alias `_rowoffset` as `__lancedb_take_row_offset`. if isinstance(requested_columns, dict): - requested_columns = list(requested_columns) + projection = requested_columns + else: + projection = {column: column for column in requested_columns} - data = {} - for column in requested_columns: - if column == "a": - data[column] = [REMOTE_ROWS[offset] for offset in offsets] - elif column == "_rowoffset": - data[column] = offsets - elif column == "_rowid": - data[column] = offsets + sources = { + "a": [REMOTE_ROWS[offset] for offset in offsets], + "_rowoffset": offsets, + "_rowid": offsets, + } + data = { + name: sources[source] + for name, source in projection.items() + if source in sources + } table = pa.table(data) request.send_response(200)