From 96a3c36682358d6b9dba2b4bf135e7e03db1fbfa Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 20 Aug 2026 15:55:18 +0000 Subject: [PATCH 1/2] fix(kernel): pass cursor row limit to kernel --- KERNEL_REV | 2 +- src/databricks/sql/backend/kernel/client.py | 11 +-- .../sql/backend/kernel/result_set.py | 72 +++++++++---------- tests/e2e/test_kernel_backend.py | 5 +- tests/unit/test_kernel_client.py | 13 ++-- tests/unit/test_kernel_result_set.py | 62 +--------------- 6 files changed, 46 insertions(+), 119 deletions(-) diff --git a/KERNEL_REV b/KERNEL_REV index 97019339d..f74d8d55b 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -45a0d6ae1de2f203220913ba96c994ebb2d7aae4 +9e3dbf9c40733b176151e001c9a15202030b967a diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 8df7e887d..eeb496f73 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -484,6 +484,7 @@ def execute_command( try: try: stmt.set_sql(operation) + stmt.set_row_limit(row_limit) if query_tags: # Per-statement query tags. The kernel serialises the # dict (None value -> bare key) into the SEA @@ -590,9 +591,7 @@ 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, row_limit=row_limit - ) + return self._make_result_set(executed, cursor, command_id) except Exception as exc: raise _wrap_kernel_exception("execute_command", exc) from exc @@ -764,9 +763,7 @@ 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, row_limit=cursor.row_limit - ) + return self._make_result_set(stream, cursor, command_id) except Exception as exc: raise _wrap_kernel_exception("get_execution_result", exc) from exc @@ -777,7 +774,6 @@ 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 @@ -799,7 +795,6 @@ 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 290df1de2..ed98984c8 100644 --- a/src/databricks/sql/backend/kernel/result_set.py +++ b/src/databricks/sql/backend/kernel/result_set.py @@ -21,9 +21,6 @@ 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 @@ -70,7 +67,6 @@ def __init__( command_id: CommandId, arraysize: int, buffer_size_bytes: int, - row_limit: Optional[int] = None, ): try: schema = kernel_handle.arrow_schema() @@ -104,55 +100,27 @@ 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._mark_exhausted() + self._exhausted = True + self.has_more_rows = False + self.status = CommandState.SUCCEEDED 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: @@ -188,10 +156,36 @@ def _take_buffered(self, n: int) -> pyarrow.Table: return pyarrow.Table.from_batches(slices, schema=self._schema) def _drain(self) -> pyarrow.Table: - """Consume the remaining logical stream into one table.""" - while not self._exhausted: - self._pull_one_batch() - return self._take_buffered(self._buffered_count) + """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) # ----- Arrow fetches ----- diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py index 55115f037..9b9b4487c 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/tests/e2e/test_kernel_backend.py @@ -183,11 +183,12 @@ def test_fetchall_arrow(conn): assert table.column_names == ["a", "b"] -@pytest.mark.parametrize("row_limit", [0, 1, 5]) +@pytest.mark.parametrize("row_limit", [None, 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)) + expected = list(range(10 if row_limit is None else row_limit)) + assert [row[0] for row in cur.fetchall()] == expected # ─── Logging (Rust kernel -> Python logging bridge) ────────────────────────── diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index ca595d8f2..29290088d 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -436,7 +436,8 @@ def test_execute_command_forwards_query_tags(): assert stmt.execute.called -def test_execute_command_applies_row_limit_to_result_set(): +@pytest.mark.parametrize("row_limit", [None, 1, 5]) +def test_execute_command_forwards_row_limit(row_limit): c = _make_client() c._kernel_session = MagicMock() cursor = MagicMock() @@ -450,7 +451,7 @@ def test_execute_command_applies_row_limit_to_result_set(): ) c._kernel_session.statement.return_value = stmt - result = c.execute_command( + c.execute_command( operation="SELECT * FROM range(10)", session_id=MagicMock(), max_rows=1, @@ -461,11 +462,10 @@ def test_execute_command_applies_row_limit_to_result_set(): parameters=[], async_op=False, enforce_embedded_schema_correctness=False, - row_limit=5, + row_limit=row_limit, ) - assert result is not None - assert result._row_limit == 5 + stmt.set_row_limit.assert_called_once_with(row_limit) # --------------------------------------------------------------------------- @@ -809,13 +809,11 @@ 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() @@ -1067,7 +1065,6 @@ 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 fe93e4e85..9ec69380a 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/tests/unit/test_kernel_result_set.py @@ -28,7 +28,6 @@ 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 @@ -36,7 +35,6 @@ 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() @@ -45,7 +43,7 @@ def close(self): self.closed = True -def _make_rs(handle, row_limit=None) -> KernelResultSet: +def _make_rs(handle) -> KernelResultSet: # The base ResultSet __init__ takes a `connection` ref it never # actually dereferences during these buffer tests, so a Mock is # fine. @@ -58,7 +56,6 @@ def _make_rs(handle, row_limit=None) -> KernelResultSet: command_id=CommandId.from_sea_statement_id("smoke-test"), arraysize=100, buffer_size_bytes=1024, - row_limit=row_limit, ) @@ -143,63 +140,6 @@ 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_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) - - 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 a846edb002dbb386e527a6d1cb6ae883e5ec11ff Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 20 Aug 2026 17:13:38 +0000 Subject: [PATCH 2/2] test(kernel): cover zero row limit --- tests/e2e/test_kernel_backend.py | 2 +- tests/unit/test_kernel_client.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py index 9b9b4487c..700528c82 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/tests/e2e/test_kernel_backend.py @@ -183,7 +183,7 @@ def test_fetchall_arrow(conn): assert table.column_names == ["a", "b"] -@pytest.mark.parametrize("row_limit", [None, 1, 5]) +@pytest.mark.parametrize("row_limit", [None, 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") diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 29290088d..3253c6fff 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -436,7 +436,7 @@ def test_execute_command_forwards_query_tags(): assert stmt.execute.called -@pytest.mark.parametrize("row_limit", [None, 1, 5]) +@pytest.mark.parametrize("row_limit", [None, 0, 1, 5]) def test_execute_command_forwards_row_limit(row_limit): c = _make_client() c._kernel_session = MagicMock()