Fix evaluate limit bug (#388)

* Fix evaluate limit bug

* Only use model LIMIT if no limit is provided in evaluate

* Fix import

* Fix import

* PR feedback

Author: georgesittas

sqlmesh/core/constants.py

@@ -28,3 +28,6 @@
 
 JINJA_MACROS = "JINJA_MACROS"
 """Used to store jinja macros in the execution env."""
+
+DEFAULT_MAX_LIMIT = 1000
+"""The default maximum row limit that is used when evaluating a model."""

sqlmesh/core/context.py

@@ -42,6 +42,7 @@
 
 import pandas as pd
 from sqlglot import exp
+from sqlglot.executor import execute
 
 from sqlmesh.core import constants as c
 from sqlmesh.core._typing import NotificationTarget
@@ -55,7 +56,7 @@
 from sqlmesh.core.hooks import hook
 from sqlmesh.core.loader import Loader, SqlMeshLoader, update_model_schemas
 from sqlmesh.core.macros import ExecutableOrMacro
-from sqlmesh.core.model import Model
+from sqlmesh.core.model import Model, SqlModel
 from sqlmesh.core.plan import Plan
 from sqlmesh.core.scheduler import Scheduler
 from sqlmesh.core.snapshot import (
@@ -521,7 +522,10 @@ def evaluate(
             start: The start of the interval to evaluate.
             end: The end of the interval to evaluate.
             latest: The latest time used for non incremental datasets.
-            limit: A limit applied to the model, this must be > 0.
+            limit: A limit applied to the model, this must be > 0. If this argument is omitted
+                and the model contains a LIMIT clause, then the clause's expression will be used
+                instead. In any case, the final limit is bounded by the DEFAULT_MAX_LIMIT constant.
+                Default: DEFAULT_MAX_LIMIT
         """
         if isinstance(model_or_snapshot, str):
             snapshot = self.snapshots[model_or_snapshot]
@@ -531,15 +535,20 @@ def evaluate(
             snapshot = self.snapshots[model_or_snapshot.name]
 
         if not limit or limit <= 0:
-            limit = 1000
+            limit = c.DEFAULT_MAX_LIMIT
+            if isinstance(snapshot.model, SqlModel):
+                model_limit = snapshot.model.render_query().args.get("limit")
+
+                if model_limit:
+                    limit = min(limit, execute(exp.select(model_limit.expression)).rows[0][0])
 
         df = self.snapshot_evaluator.evaluate(
             snapshot,
             start,
             end,
             latest,
             snapshots=self.snapshots,
-            limit=limit,
+            limit=t.cast(int, limit),
         )
 
         if df is None:

tests/core/test_context.py

@@ -8,6 +8,8 @@
 import sqlmesh.core.constants
 from sqlmesh.core.config import Config, ModelDefaultsConfig
 from sqlmesh.core.context import Context
+from sqlmesh.core.dialect import parse
+from sqlmesh.core.model import load_model
 from sqlmesh.core.plan import BuiltInPlanEvaluator, Plan
 from sqlmesh.utils.errors import ConfigError
 from tests.utils.test_filesystem import create_temp_file
@@ -234,6 +236,24 @@ def test_diff(sushi_context: Context, mocker: MockerFixture):
     assert mock_console.show_model_difference_summary.called
 
 
+def test_evaluate_limit():
+    context = Context(config=Config())
+
+    context.upsert_model(
+        load_model(
+            parse(
+                """
+        MODEL(name limit_test, kind FULL);
+        SELECT t.v as v FROM (VALUES (1), (2), (3), (4), (5)) AS t(v) LIMIT 1 + 2"""
+            )
+        )
+    )
+
+    assert context.evaluate("limit_test", "2020-01-01", "2020-01-02", "2020-01-02").size == 3
+    assert context.evaluate("limit_test", "2020-01-01", "2020-01-02", "2020-01-02", 4).size == 4
+    assert context.evaluate("limit_test", "2020-01-01", "2020-01-02", "2020-01-02", 2).size == 2
+
+
 def test_ignore_files(mocker: MockerFixture, tmpdir):
     mocker.patch.object(
         sqlmesh.core.constants,