feat(nodejs): add order_by method to Query (#3123)

python/python/tests/test_fts.py

@@ -29,6 +29,7 @@ from lancedb.query import (
     MultiMatchQuery,
     PhraseQuery,
     BooleanQuery,
+    ColumnOrdering,
     Occur,
     LanceFtsQueryBuilder,
 )
@@ -499,6 +500,36 @@ async def test_search_fts_specify_column_async(async_table):
         pass
 
 
+def test_search_order_by_descending(table):
+    table.create_fts_index("text")
+    rows = (
+        table.search("puppy")
+        .order_by([ColumnOrdering(column_name="count", ascending=False)])
+        .limit(20)
+        .select(["text", "count"])
+        .to_list()
+    )
+
+    for r in rows:
+        assert "puppy" in r["text"]
+    assert sorted(rows, key=lambda x: x["count"], reverse=True) == rows
+
+
+def test_search_order_by_ascending(table):
+    table.create_fts_index("text")
+    rows = (
+        table.search("puppy")
+        .order_by([ColumnOrdering(column_name="count", ascending=True)])
+        .limit(20)
+        .select(["text", "count"])
+        .to_list()
+    )
+
+    for r in rows:
+        assert "puppy" in r["text"]
+    assert sorted(rows, key=lambda x: x["count"]) == rows
+
+
 def test_create_index_from_table(tmp_path, table):
     table.create_fts_index("text")
     df = table.search("puppy").limit(5).select(["text"]).to_pandas()

python/python/tests/test_query.py

@@ -25,6 +25,7 @@ from lancedb.query import (
     AsyncHybridQuery,
     AsyncQueryBase,
     AsyncVectorQuery,
+    ColumnOrdering,
     LanceVectorQueryBuilder,
     MatchQuery,
     PhraseQuery,
@@ -164,6 +165,71 @@ def test_offset(table):
     assert len(results_with_offset.to_pandas()) == 1
 
 
+def test_order_by_plain_query(mem_db):
+    table = mem_db.create_table(
+        "test_order_by",
+        pa.table(
+            {
+                "group": [1, 1, 1, 2],
+                "score": [None, 1.0, 1.0, 0.5],
+                "name": ["z", "b", "a", "c"],
+            }
+        ),
+    )
+
+    res = (
+        table.search()
+        .order_by(
+            [
+                ColumnOrdering(column_name="group", ascending=True, nulls_first=False),
+                ColumnOrdering(column_name="score", ascending=True, nulls_first=True),
+                ColumnOrdering(column_name="name", ascending=True, nulls_first=False),
+            ]
+        )
+        .to_arrow()
+    )
+
+    assert res.select(["group", "score", "name"]).to_pylist() == [
+        {"group": 1, "score": None, "name": "z"},
+        {"group": 1, "score": 1.0, "name": "a"},
+        {"group": 1, "score": 1.0, "name": "b"},
+        {"group": 2, "score": 0.5, "name": "c"},
+    ]
+
+
+@pytest.mark.asyncio
+async def test_order_by_async_query(mem_db_async: AsyncConnection):
+    table = await mem_db_async.create_table(
+        "test_order_by_async",
+        pa.table(
+            {
+                "group": [1, 1, 1, 2],
+                "score": [None, 1.0, 1.0, 0.5],
+                "name": ["z", "b", "a", "c"],
+            }
+        ),
+    )
+
+    res = await (
+        table.query()
+        .order_by(
+            [
+                ColumnOrdering(column_name="group", ascending=True, nulls_first=False),
+                ColumnOrdering(column_name="score", ascending=True, nulls_first=True),
+                ColumnOrdering(column_name="name", ascending=True, nulls_first=False),
+            ]
+        )
+        .to_arrow()
+    )
+
+    assert res.select(["group", "score", "name"]).to_pylist() == [
+        {"group": 1, "score": None, "name": "z"},
+        {"group": 1, "score": 1.0, "name": "a"},
+        {"group": 1, "score": 1.0, "name": "b"},
+        {"group": 2, "score": 0.5, "name": "c"},
+    ]
+
+
 def test_query_builder(table):
     rs = (
         LanceVectorQueryBuilder(table, [0, 0], "vector")

python/python/tests/test_remote_db.py

@@ -16,6 +16,7 @@ from packaging.version import Version
 
 import lancedb
 from lancedb.conftest import MockTextEmbeddingFunction
+from lancedb.query import ColumnOrdering
 from lancedb.remote import ClientConfig
 from lancedb.remote.errors import HttpError, RetryError
 import pytest
@@ -660,6 +661,18 @@ def test_query_sync_maximal():
             "ef": None,
             "filter": "id > 0",
             "columns": ["id", "name"],
+            "order_by": [
+                {
+                    "column_name": "score",
+                    "ascending": False,
+                    "nulls_first": True,
+                },
+                {
+                    "column_name": "id",
+                    "ascending": True,
+                    "nulls_first": False,
+                },
+            ],
             "vector_column": "vector2",
             "fast_search": True,
             "with_row_id": True,
@@ -677,6 +690,14 @@ def test_query_sync_maximal():
             .refine_factor(10)
             .nprobes(5)
             .where("id > 0", prefilter=True)
+            .order_by(
+                [
+                    ColumnOrdering(
+                        column_name="score", ascending=False, nulls_first=True
+                    ),
+                    ColumnOrdering(column_name="id", ascending=True, nulls_first=False),
+                ]
+            )
             .with_row_id(True)
             .select(["id", "name"])
             .to_list()