From 54b43e2860a9c4f463f4b557bb7e962c4993bf35 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Tue, 18 Aug 2026 05:20:32 +0000 Subject: [PATCH 1/2] fix(kernel): honor cursor row limit --- src/databricks/sql/backend/kernel/client.py | 10 ++- .../sql/backend/kernel/result_set.py | 72 ++++++++++--------- tests/e2e/test_kernel_backend.py | 7 ++ tests/unit/test_kernel_client.py | 35 +++++++++ tests/unit/test_kernel_result_set.py | 54 +++++++++++++- 5 files changed, 142 insertions(+), 36 deletions(-) diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index b1a1d5b3e..8df7e887d 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/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: diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py index ed98984c8..290df1de2 100644 --- a/src/databricks/sql/backend/kernel/result_set.py +++ b/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 ----- diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py index 8b532a56a..55115f037 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/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 diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 79be53e64..ca595d8f2 100644 --- a/tests/unit/test_kernel_client.py +++ b/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) diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py index 9ec69380a..13e9bdf1b 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/tests/unit/test_kernel_result_set.py @@ -28,6 +28,7 @@ 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 @@ -35,6 +36,7 @@ def arrow_schema(self) -> pa.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): From 07fe3c9541891470ea6ca5e6ea6d4c0e38082fb4 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Tue, 18 Aug 2026 05:42:29 +0000 Subject: [PATCH 2/2] test(kernel): cover exact row limit boundary --- tests/unit/test_kernel_result_set.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py index 13e9bdf1b..fe93e4e85 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/tests/unit/test_kernel_result_set.py @@ -185,6 +185,14 @@ def test_row_limit_stops_before_fetching_extra_batches(int_schema): assert handle.fetch_calls == 1 +def test_row_limit_exact_batch_boundary_skips_exhaustion_fetch(int_schema): + handle = _FakeKernelHandle(int_schema, [_batch(int_schema, [0, 1, 2])]) + rs = _make_rs(handle, row_limit=3) + + assert rs.fetchall_arrow().column(0).to_pylist() == [0, 1, 2] + 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)