fix(kernel): honor cursor row limit

Author: vuanhphung

src/databricks/sql/backend/kernel/client.py

@@ -590,7 +590,9 @@ def execute_command(
         # native exception) — wrap the construction so callers see a
         # mapped PEP 249 exception.
         try:
-            return self._make_result_set(executed, cursor, command_id)
+            return self._make_result_set(
+                executed, cursor, command_id, row_limit=row_limit
+            )
         except Exception as exc:
             raise _wrap_kernel_exception("execute_command", exc) from exc
 
@@ -762,7 +764,9 @@ def get_execution_result(
         # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
         # can raise — map that to PEP 249 too.
         try:
-            return self._make_result_set(stream, cursor, command_id)
+            return self._make_result_set(
+                stream, cursor, command_id, row_limit=cursor.row_limit
+            )
         except Exception as exc:
             raise _wrap_kernel_exception("get_execution_result", exc) from exc
 
@@ -773,6 +777,7 @@ def _make_result_set(
         kernel_handle: Any,
         cursor: "Cursor",
         command_id: CommandId,
+        row_limit: Optional[int] = None,
     ) -> "ResultSet":
         """Build a ``KernelResultSet`` from any kernel handle. Used
         by sync execute, ``get_execution_result``, and all metadata
@@ -794,6 +799,7 @@ def _make_result_set(
             command_id=command_id,
             arraysize=cursor.arraysize,
             buffer_size_bytes=cursor.buffer_size_bytes,
+            row_limit=row_limit,
         )
 
     def _synthetic_command_id(self) -> CommandId:

src/databricks/sql/backend/kernel/result_set.py

@@ -21,6 +21,9 @@
 within a batch when ``n`` is smaller than the kernel's natural
 batch size; ``fetchall`` drains the whole stream.
 
+When a cursor has ``row_limit`` set, this class caps the logical stream
+before rows reach any of the row or Arrow fetch APIs.
+
 Note: ``buffer_size_bytes`` is accepted by the constructor for
 contract compatibility with the base ``ResultSet`` but is not
 consulted — the kernel backend currently caps buffering by rows
@@ -67,6 +70,7 @@ def __init__(
         command_id: CommandId,
         arraysize: int,
         buffer_size_bytes: int,
+        row_limit: Optional[int] = None,
     ):
         try:
             schema = kernel_handle.arrow_schema()
@@ -100,27 +104,55 @@ def __init__(
         # stays O(1) instead of walking the deque.
         self._buffered_count: int = 0
         self._exhausted: bool = False
+        # The PyO3 kernel surface does not currently expose the core
+        # StatementSpec row_limit setter. Enforce the cursor contract at
+        # this streaming boundary until it does. Negative values retain the
+        # existing unlimited behaviour; zero is a real zero-row limit.
+        self._row_limit: Optional[int] = (
+            row_limit if row_limit is not None and row_limit >= 0 else None
+        )
+        if self._row_limit == 0:
+            self._mark_exhausted()
 
     # ----- internal helpers -----
 
+    def _mark_exhausted(self) -> None:
+        self._exhausted = True
+        self.has_more_rows = False
+        self.status = CommandState.SUCCEEDED
+
+    def _remaining_row_limit(self) -> Optional[int]:
+        if self._row_limit is None:
+            return None
+        return max(
+            0,
+            self._row_limit - self._next_row_index - self._buffered_count,
+        )
+
     def _pull_one_batch(self) -> bool:
         """Pull the next batch from the kernel into the local buffer.
         Returns True if a batch was added; False if the kernel side
         is exhausted."""
         if self._exhausted:
             return False
+        remaining_limit = self._remaining_row_limit()
+        if remaining_limit == 0:
+            self._mark_exhausted()
+            return False
         try:
             batch = self._kernel_handle.fetch_next_batch()
         except Exception as exc:
             raise wrap_kernel_exception("fetch_next_batch", exc) from exc
         if batch is None:
-            self._exhausted = True
-            self.has_more_rows = False
-            self.status = CommandState.SUCCEEDED
+            self._mark_exhausted()
             return False
+        if remaining_limit is not None and batch.num_rows > remaining_limit:
+            batch = batch.slice(0, remaining_limit)
         if batch.num_rows > 0:
             self._buffer.append(batch)
             self._buffered_count += batch.num_rows
+        if remaining_limit is not None and batch.num_rows >= remaining_limit:
+            self._mark_exhausted()
         return True
 
     def _ensure_buffered(self, n_rows: int) -> int:
@@ -156,36 +188,10 @@ def _take_buffered(self, n: int) -> pyarrow.Table:
         return pyarrow.Table.from_batches(slices, schema=self._schema)
 
     def _drain(self) -> pyarrow.Table:
-        """Consume everything left in the buffer + kernel stream
-        and return as a single Table."""
-        chunks: List[pyarrow.RecordBatch] = []
-        if self._buffer and self._buffer_offset > 0:
-            head = self._buffer.popleft()
-            chunks.append(
-                head.slice(self._buffer_offset, head.num_rows - self._buffer_offset)
-            )
-            self._buffer_offset = 0
-        while self._buffer:
-            chunks.append(self._buffer.popleft())
-        if not self._exhausted:
-            while True:
-                try:
-                    batch = self._kernel_handle.fetch_next_batch()
-                except Exception as exc:
-                    raise wrap_kernel_exception("fetch_next_batch", exc) from exc
-                if batch is None:
-                    self._exhausted = True
-                    self.has_more_rows = False
-                    self.status = CommandState.SUCCEEDED
-                    break
-                if batch.num_rows > 0:
-                    chunks.append(batch)
-        rows = sum(c.num_rows for c in chunks)
-        self._buffered_count = 0
-        self._next_row_index += rows
-        if not chunks:
-            return pyarrow.Table.from_batches([], schema=self._schema)
-        return pyarrow.Table.from_batches(chunks, schema=self._schema)
+        """Consume the remaining logical stream into one table."""
+        while not self._exhausted:
+            self._pull_one_batch()
+        return self._take_buffered(self._buffered_count)
 
     # ----- Arrow fetches -----
 

tests/e2e/test_kernel_backend.py

@@ -183,6 +183,13 @@ def test_fetchall_arrow(conn):
         assert table.column_names == ["a", "b"]
 
 
+@pytest.mark.parametrize("row_limit", [0, 1, 5])
+def test_cursor_row_limit(conn, row_limit):
+    with conn.cursor(row_limit=row_limit) as cur:
+        cur.execute("SELECT id FROM range(10) ORDER BY id")
+        assert [row[0] for row in cur.fetchall()] == list(range(row_limit))
+
+
 # ─── Logging (Rust kernel -> Python logging bridge) ──────────────────────────
 #
 # Layer 3 of the logger-name drift guard (see also the Rust tests

tests/unit/test_kernel_client.py

@@ -436,6 +436,38 @@ def test_execute_command_forwards_query_tags():
     assert stmt.execute.called
 
 
+def test_execute_command_applies_row_limit_to_result_set():
+    c = _make_client()
+    c._kernel_session = MagicMock()
+    cursor = MagicMock()
+    cursor.arraysize = 100
+    cursor.buffer_size_bytes = 1024
+
+    stmt = MagicMock()
+    stmt.execute.return_value = MagicMock(
+        statement_id="stmt-id",
+        arrow_schema=MagicMock(return_value=pa.schema([("x", pa.int64())])),
+    )
+    c._kernel_session.statement.return_value = stmt
+
+    result = c.execute_command(
+        operation="SELECT * FROM range(10)",
+        session_id=MagicMock(),
+        max_rows=1,
+        max_bytes=1,
+        lz4_compression=False,
+        cursor=cursor,
+        use_cloud_fetch=False,
+        parameters=[],
+        async_op=False,
+        enforce_embedded_schema_correctness=False,
+        row_limit=5,
+    )
+
+    assert result is not None
+    assert result._row_limit == 5
+
+
 # ---------------------------------------------------------------------------
 # Staging / volume operations — fail loud (not silently no-op)
 # ---------------------------------------------------------------------------
@@ -777,11 +809,13 @@ def test_get_execution_result_attaches_by_id():
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
+    cursor.row_limit = 5
     cid = CommandId.from_sea_statement_id("async-1")
 
     rs = c.get_execution_result(cid, cursor=cursor)
 
     assert rs is not None
+    assert rs._row_limit == 5
     c._kernel_session.attach_async_statement.assert_called_with("async-1")
     handle.await_result.assert_called_once_with()
 
@@ -1033,6 +1067,7 @@ def test_get_execution_result_is_re_callable():
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
+    cursor.row_limit = None
 
     rs1 = c.get_execution_result(cid, cursor=cursor)
     rs2 = c.get_execution_result(cid, cursor=cursor)

tests/unit/test_kernel_result_set.py

@@ -28,13 +28,15 @@ def __init__(self, schema: pa.Schema, batches):
         self._schema = schema
         self._batches: Deque[pa.RecordBatch] = deque(batches)
         self.closed = False
+        self.fetch_calls = 0
 
     def arrow_schema(self) -> pa.Schema:
         return self._schema
 
     def fetch_next_batch(self):
         if self.closed:
             raise RuntimeError("fetched after close")
+        self.fetch_calls += 1
         if not self._batches:
             return None
         return self._batches.popleft()
@@ -43,7 +45,7 @@ def close(self):
         self.closed = True
 
 
-def _make_rs(handle) -> KernelResultSet:
+def _make_rs(handle, row_limit=None) -> KernelResultSet:
     # The base ResultSet __init__ takes a `connection` ref it never
     # actually dereferences during these buffer tests, so a Mock is
     # fine.
@@ -56,6 +58,7 @@ def _make_rs(handle) -> KernelResultSet:
         command_id=CommandId.from_sea_statement_id("smoke-test"),
         arraysize=100,
         buffer_size_bytes=1024,
+        row_limit=row_limit,
     )
 
 
@@ -140,6 +143,55 @@ def test_fetchall_rows(int_schema):
     assert [r[0] for r in rows] == [1, 2, 3]
 
 
+@pytest.mark.parametrize("row_limit", [0, 1, 5])
+def test_row_limit_caps_fetchall(int_schema, row_limit):
+    handle = _FakeKernelHandle(
+        int_schema,
+        [_batch(int_schema, [0, 1, 2]), _batch(int_schema, list(range(3, 10)))],
+    )
+    rs = _make_rs(handle, row_limit=row_limit)
+
+    rows = rs.fetchall()
+
+    assert [row[0] for row in rows] == list(range(row_limit))
+    assert rs.rownumber == row_limit
+
+
+def test_row_limit_applies_across_fetch_methods(int_schema):
+    handle = _FakeKernelHandle(
+        int_schema,
+        [_batch(int_schema, [0, 1, 2]), _batch(int_schema, [3, 4, 5, 6])],
+    )
+    rs = _make_rs(handle, row_limit=5)
+
+    first = rs.fetchmany(2)
+    third = rs.fetchone()
+    rest = rs.fetchall_arrow()
+
+    assert [row[0] for row in first] == [0, 1]
+    assert third is not None and third[0] == 2
+    assert rest.column(0).to_pylist() == [3, 4]
+    assert rs.fetchone() is None
+
+
+def test_row_limit_stops_before_fetching_extra_batches(int_schema):
+    handle = _FakeKernelHandle(
+        int_schema,
+        [_batch(int_schema, [0, 1, 2]), _batch(int_schema, [3, 4, 5])],
+    )
+    rs = _make_rs(handle, row_limit=2)
+
+    assert rs.fetchall_arrow().column(0).to_pylist() == [0, 1]
+    assert handle.fetch_calls == 1
+
+
+def test_row_limit_larger_than_result_returns_all_rows(int_schema):
+    handle = _FakeKernelHandle(int_schema, [_batch(int_schema, [1, 2, 3])])
+    rs = _make_rs(handle, row_limit=10)
+
+    assert rs.fetchall_arrow().column(0).to_pylist() == [1, 2, 3]
+
+
 def test_fetchmany_negative_raises(int_schema):
     rs = _make_rs(_FakeKernelHandle(int_schema, []))
     with pytest.raises(ValueError):