Repository: huggingface/dataset-viewer

Issue:
Sort the rows along a column
See https://github.com/huggingface/moon-landing/issues/9201 (internal)




Discussion:
Cc @albertvillanova since it's quite similar to filtering
> Cc @albertvillanova since it's quite similar to filtering

by the way: should it be served by /filter?
Or: by a new endpoint that takes a `?sql=` parameter? Maybe restricted to calls from the Hub at the beginning.

This endpoint might even be `/rows`, and we could handle all the cases (/first-rows, random access, filter, search, sort, SQL) in one big endpoint.

Resolving patch:
diff --git a/services/search/src/search/routes/filter.py b/services/search/src/search/routes/filter.py
--- a/services/search/src/search/routes/filter.py
+++ b/services/search/src/search/routes/filter.py
@@ -42,14 +42,15 @@
 FILTER_QUERY = """\
     SELECT {columns}
     FROM data
-    WHERE {where}
+    {where}
+    {orderby}
     LIMIT {limit}
     OFFSET {offset}"""
 
 FILTER_COUNT_QUERY = """\
     SELECT COUNT(*)
     FROM data
-    WHERE {where}"""
+    {where}"""
 
 SQL_INVALID_SYMBOLS = "|".join([";", "--", r"/\*", r"\*/"])
 SQL_INVALID_SYMBOLS_PATTERN = re.compile(rf"(?:{SQL_INVALID_SYMBOLS})", flags=re.IGNORECASE)
@@ -83,11 +84,15 @@ async def filter_endpoint(request: Request) -> Response:
                     dataset = get_request_parameter(request, "dataset", required=True)
                     config = get_request_parameter(request, "config", required=True)
                     split = get_request_parameter(request, "split", required=True)
-                    where = get_request_parameter(request, "where", required=True)
-                    validate_where_parameter(where)
+                    where = get_request_parameter(request, "where")
+                    validate_query_parameter(where, "where")
+                    orderby = get_request_parameter(request, "orderby")
+                    validate_query_parameter(orderby, "orderby")
                     offset = get_request_parameter_offset(request)
                     length = get_request_parameter_length(request)
-                    logger.info(f"/filter, {dataset=}, {config=}, {split=}, {where=}, {offset=}, {length=}")
+                    logger.info(
+                        f"/filter, {dataset=}, {config=}, {split=}, {where=}, {orderby=}, {offset=}, {length=}"
+                    )
                 with StepProfiler(method="filter_endpoint", step="check authentication"):
                     # If auth_check fails, it will raise an exception that will be caught below
                     await auth_check(
@@ -154,6 +159,7 @@ async def filter_endpoint(request: Request) -> Response:
                         index_file_location,
                         supported_columns,
                         where,
+                        orderby,
                         length,
                         offset,
                         extensions_directory,
@@ -195,6 +201,7 @@ def execute_filter_query(
     index_file_location: str,
     columns: list[str],
     where: str,
+    orderby: str,
     limit: int,
     offset: int,
     extensions_directory: Optional[str] = None,
@@ -202,19 +209,20 @@ def execute_filter_query(
     with duckdb_connect(extensions_directory=extensions_directory, database=index_file_location) as con:
         filter_query = FILTER_QUERY.format(
             columns=",".join([f'"{column}"' for column in columns]),
-            where=where,
+            where=f"WHERE {where}" if where else "",
+            orderby=f"ORDER BY {orderby}" if orderby else "",
             limit=limit,
             offset=offset,
         )
-        filter_count_query = FILTER_COUNT_QUERY.format(where=where)
+        filter_count_query = FILTER_COUNT_QUERY.format(where=f"WHERE {where}" if where else "")
         try:
             pa_table = con.sql(filter_query).arrow()
             num_rows_total = con.sql(filter_count_query).fetchall()[0][0]
-        except duckdb.Error:
-            raise InvalidParameterError(message="Parameter 'where' is invalid")
+        except duckdb.Error as err:
+            raise InvalidParameterError(message="A query parameter is invalid") from err
     return num_rows_total, pa_table
 
 
-def validate_where_parameter(where: str) -> None:
-    if SQL_INVALID_SYMBOLS_PATTERN.search(where):
-        raise InvalidParameterError(message="Parameter 'where' contains invalid symbols")
+def validate_query_parameter(parameter_value: str, parameter_name: str) -> None:
+    if SQL_INVALID_SYMBOLS_PATTERN.search(parameter_value):
+        raise InvalidParameterError(message=f"Parameter '{parameter_name}' contains invalid symbols")


Test patch:
diff --git a/e2e/tests/test_53_filter.py b/e2e/tests/test_53_filter.py
--- a/e2e/tests/test_53_filter.py
+++ b/e2e/tests/test_53_filter.py
@@ -13,9 +13,10 @@ def test_filter_endpoint(normal_user_public_dataset: str) -> None:
     offset = 1
     length = 2
     where = "col_4='B'"
+    orderby = "col_2 DESC"
     filter_response = poll_until_ready_and_assert(
         relative_url=(
-            f"/filter?dataset={dataset}&config={config}&split={split}&offset={offset}&length={length}&where={where}"
+            f"/filter?dataset={dataset}&config={config}&split={split}&offset={offset}&length={length}&where={where}&orderby={orderby}"
         ),
         check_x_revision=True,
         dataset=dataset,
@@ -44,13 +45,13 @@ def test_filter_endpoint(normal_user_public_dataset: str) -> None:
         "truncated_cells": [],
     }, rows[0]
     assert rows[1] == {
-        "row_idx": 3,
+        "row_idx": 0,
         "row": {
-            "col_1": "The wingman spots the pirateship coming at him and warns the Dark Lord",
-            "col_2": 3,
-            "col_3": 3.0,
+            "col_1": "There goes another one.",
+            "col_2": 0,
+            "col_3": 0.0,
             "col_4": "B",
-            "col_5": None,
+            "col_5": True,
         },
         "truncated_cells": [],
     }, rows[1]
@@ -66,6 +67,7 @@ def test_filter_endpoint(normal_user_public_dataset: str) -> None:
 @pytest.mark.parametrize(
     "where,expected_num_rows",
     [
+        ("", 4),
         ("col_2=3", 1),
         ("col_2<3", 3),
         ("col_2>3", 0),
@@ -79,8 +81,11 @@ def test_filter_endpoint(normal_user_public_dataset: str) -> None:
 def test_filter_endpoint_parameter_where(where: str, expected_num_rows: int, normal_user_public_dataset: str) -> None:
     dataset = normal_user_public_dataset
     config, split = get_default_config_split()
+    relative_url = f"/filter?dataset={dataset}&config={config}&split={split}"
+    if where:
+        relative_url += f"&where={where}"
     response = poll_until_ready_and_assert(
-        relative_url=f"/filter?dataset={dataset}&config={config}&split={split}&where={where}",
+        relative_url=relative_url,
         check_x_revision=True,
         dataset=dataset,
     )
@@ -89,6 +94,34 @@ def test_filter_endpoint_parameter_where(where: str, expected_num_rows: int, nor
     assert len(content["rows"]) == expected_num_rows
 
 
+@pytest.mark.parametrize(
+    "orderby, expected_first_row_idx",
+    [
+        ("", 1),
+        ("col_4", 2),
+        ("col_3 DESC", 3),
+    ],
+)
+def test_filter_endpoint_parameter_orderby(
+    orderby: str, expected_first_row_idx: int, normal_user_public_dataset: str
+) -> None:
+    dataset = normal_user_public_dataset
+    config, split = get_default_config_split()
+    where = "col_2>0"
+    relative_url = f"/filter?dataset={dataset}&config={config}&split={split}&where={where}"
+    if orderby:
+        relative_url += f"&orderby={orderby}"
+    response = poll_until_ready_and_assert(
+        relative_url=relative_url,
+        check_x_revision=True,
+        dataset=dataset,
+    )
+    content = response.json()
+    assert "rows" in content, response
+    rows = content["rows"]
+    assert rows[0]["row_idx"] == expected_first_row_idx, rows[0]
+
+
 def test_filter_images_endpoint(normal_user_images_public_dataset: str) -> None:
     dataset = normal_user_images_public_dataset
     config, split = get_default_config_split()
diff --git a/services/search/tests/routes/test_filter.py b/services/search/tests/routes/test_filter.py
--- a/services/search/tests/routes/test_filter.py
+++ b/services/search/tests/routes/test_filter.py
@@ -7,6 +7,7 @@
 
 import duckdb
 import pyarrow as pa
+import pyarrow.compute as pc
 import pytest
 from datasets import Dataset
 from libapi.exceptions import InvalidParameterError
@@ -14,7 +15,7 @@
 from libcommon.storage_client import StorageClient
 
 from search.config import AppConfig
-from search.routes.filter import execute_filter_query, validate_where_parameter
+from search.routes.filter import execute_filter_query, validate_query_parameter
 
 CACHED_ASSETS_FOLDER = "cached-assets"
 
@@ -60,31 +61,59 @@ def index_file_location(ds: Dataset) -> Generator[str, None, None]:
     os.remove(index_file_location)
 
 
-@pytest.mark.parametrize("where", ["col='A'"])
-def test_validate_where_parameter(where: str) -> None:
-    validate_where_parameter(where)
+@pytest.mark.parametrize(
+    "parameter_name, parameter_value", [("where", "col='A'"), ("orderby", "A"), ("orderby", "A DESC")]
+)
+def test_validate_query_parameter(parameter_name: str, parameter_value: str) -> None:
+    validate_query_parameter(parameter_value, parameter_name)
 
 
-@pytest.mark.parametrize("where", ["col='A'; SELECT * from data", "col='A' /*", "col='A'--"])
-def test_validate_where_parameter_raises(where: str) -> None:
+@pytest.mark.parametrize("sql_injection", ["; SELECT * from data", " /*", "--"])
+@pytest.mark.parametrize(
+    "parameter_name, parameter_value",
+    [("where", "col='A'"), ("orderby", "A"), ("orderby", "A DESC")],
+)
+def test_validate_query_parameter_raises(parameter_name: str, parameter_value: str, sql_injection: str) -> None:
     with pytest.raises(InvalidParameterError):
-        validate_where_parameter(where)
+        validate_query_parameter(parameter_value + sql_injection, parameter_name)
 
 
+@pytest.mark.parametrize("orderby", ["", "age", "age DESC"])
+@pytest.mark.parametrize("where", ["", "gender='female'"])
 @pytest.mark.parametrize("columns", [["name", "age"], ["name"]])
-def test_execute_filter_query(index_file_location: str, columns: list[str]) -> None:
+def test_execute_filter_query(columns: list[str], where: str, orderby: str, index_file_location: str) -> None:
     # in split-duckdb-index we always add the ROW_IDX_COLUMN column
     # see https://github.com/huggingface/dataset-viewer/blob/main/services/worker/src/worker/job_runners/split/duckdb_index.py#L305
-    columns.append(ROW_IDX_COLUMN)
-    where, limit, offset = "gender='female'", 1, 1
+    columns = columns + [ROW_IDX_COLUMN]
+    limit, offset = 1, 1
     num_rows_total, pa_table = execute_filter_query(
-        index_file_location=index_file_location, columns=columns, where=where, limit=limit, offset=offset
+        index_file_location=index_file_location,
+        columns=columns,
+        where=where,
+        orderby=orderby,
+        limit=limit,
+        offset=offset,
     )
-    assert num_rows_total == 2
-    expected = pa.Table.from_pydict(
-        {"__hf_index_id": [3], "name": ["Simone"], "gender": ["female"], "age": [30]}
-    ).select(columns)
-    assert pa_table == expected
+    expected_num_rows_total = 2 if where else 4
+    assert num_rows_total == expected_num_rows_total
+    expected_pa_table = pa.Table.from_pydict(
+        {
+            "__hf_index_id": [0, 1, 2, 3],
+            "name": ["Marie", "Paul", "Leo", "Simone"],
+            "gender": ["female", "male", "male", "female"],
+            "age": [35, 30, 25, 30],
+        }
+    )
+    if where:
+        expected_pa_table = expected_pa_table.filter(pc.field("gender") == "female")
+    if orderby:
+        if orderby.endswith(" DESC"):
+            sorting = [(orderby.removesuffix(" DESC"), "descending")]
+            expected_pa_table = expected_pa_table.sort_by(sorting)
+        else:
+            expected_pa_table = expected_pa_table.sort_by(orderby)
+    expected_pa_table = expected_pa_table.slice(offset, limit).select(columns)
+    assert pa_table == expected_pa_table
 
 
 @pytest.mark.parametrize("where", ["non-existing-column=30", "name=30", "name>30"])
@@ -92,7 +121,12 @@ def test_execute_filter_query_raises(where: str, index_file_location: str) -> No
     columns, limit, offset = ["name", "gender", "age"], 100, 0
     with pytest.raises(InvalidParameterError):
         _ = execute_filter_query(
-            index_file_location=index_file_location, columns=columns, where=where, limit=limit, offset=offset
+            index_file_location=index_file_location,
+            columns=columns,
+            where=where,
+            orderby="",
+            limit=limit,
+            offset=offset,
         )