feat: sql to substrait - add support for limit and offset

Author: tokoko

src/substrait/builders/plan.py

@@ -26,8 +26,8 @@
 
 def _merge_extensions(*objs):
     return {
-        "extension_uris": merge_extension_uris(*[b.extension_uris for b in objs]),
-        "extensions": merge_extension_declarations(*[b.extensions for b in objs]),
+        "extension_uris": merge_extension_uris(*[b.extension_uris for b in objs if b]),
+        "extensions": merge_extension_declarations(*[b.extensions for b in objs if b]),
     }
 
 
@@ -193,13 +193,15 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan:
         bound_plan = plan if isinstance(plan, stp.Plan) else plan(registry)
         ns = infer_plan_schema(bound_plan)
 
-        bound_offset = resolve_expression(offset, ns, registry)
+        bound_offset = resolve_expression(offset, ns, registry) if offset else None
         bound_count = resolve_expression(count, ns, registry)
 
         rel = stalg.Rel(
             fetch=stalg.FetchRel(
                 input=bound_plan.relations[-1].root.input,
-                offset_expr=bound_offset.referred_expr[0].expression,
+                offset_expr=bound_offset.referred_expr[0].expression
+                if bound_offset
+                else None,
                 count_expr=bound_count.referred_expr[0].expression,
             )
         )

src/substrait/sql/sql_to_substrait.py

@@ -14,6 +14,7 @@
     project,
     filter,
     sort,
+    fetch,
     set,
     join,
     aggregate,
@@ -35,9 +36,7 @@
     "Eq": ("functions_comparison.yaml", "equal"),
 }
 
-aggregate_function_mapping = {
-    "SUM": ("functions_arithmetic.yaml", "sum"),
-}
+aggregate_function_mapping = {"SUM": ("functions_arithmetic.yaml", "sum")}
 
 window_function_mapping = {
     "row_number": ("functions_arithmetic.yaml", "row_number"),
@@ -190,6 +189,29 @@ def translate(ast: dict, schema_resolver: SchemaResolver, registry: ExtensionReg
                 for e in ast["order_by"]["kind"]["Expressions"]
             ]
             relation = sort(relation, expressions)(registry)
+
+        if ast["limit_clause"]:
+            limit_expression = translate_expression(
+                ast["limit_clause"]["LimitOffset"]["limit"],
+                schema_resolver=schema_resolver,
+                registry=registry,
+                measures=None,
+                groupings=None,
+            )
+
+            if ast["limit_clause"]["LimitOffset"]["offset"]:
+                offset_expression = translate_expression(
+                    ast["limit_clause"]["LimitOffset"]["offset"]["value"],
+                    schema_resolver=schema_resolver,
+                    registry=registry,
+                    measures=None,
+                    groupings=None,
+                )
+            else:
+                offset_expression = None
+
+            relation = fetch(relation, offset_expression, limit_expression)(registry)
+
         return relation
     elif op == "Select":
         relation = translate(

tests/sql/test_sql_to_substrait.py

@@ -250,6 +250,48 @@ def test_order_by(engine: str):
     )
 
 
+@pytest.mark.parametrize(
+    "engine",
+    [
+        pytest.param(
+            "duckdb",
+            marks=[
+                pytest.mark.skipif(
+                    sys.platform.startswith("win"),
+                    reason="duckdb substrait extension not found on windows",
+                ),
+                pytest.mark.xfail,
+            ],
+        ),
+        "datafusion",
+    ],
+)
+def test_select_limit(engine: str):
+    assert_query("""SELECT store_id FROM stores ORDER BY store_id LIMIT 2""", engine)
+
+
+@pytest.mark.parametrize(
+    "engine",
+    [
+        pytest.param(
+            "duckdb",
+            marks=[
+                pytest.mark.skipif(
+                    sys.platform.startswith("win"),
+                    reason="duckdb substrait extension not found on windows",
+                ),
+                pytest.mark.xfail,
+            ],
+        ),
+        "datafusion",
+    ],
+)
+def test_select_limit_offset(engine: str):
+    assert_query(
+        """SELECT store_id FROM stores ORDER BY store_id LIMIT 2 OFFSET 2""", engine
+    )
+
+
 @pytest.mark.parametrize(
     "engine",
     [