Skip to content

Commit a68464c

Browse files
vuanhphungvuanhphung
authored
fix(kernel): pass cursor row limit to kernel (#922)
* fix(kernel): pass cursor row limit to kernel * test(kernel): cover zero row limit
1 parent 0d9126e commit a68464c

6 files changed

Lines changed: 46 additions & 119 deletions

File tree

‎KERNEL_REV‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
45a0d6ae1de2f203220913ba96c994ebb2d7aae4
1+
9e3dbf9c40733b176151e001c9a15202030b967a

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

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -484,6 +484,7 @@ def execute_command(
484484
try:
485485
try:
486486
stmt.set_sql(operation)
487+
stmt.set_row_limit(row_limit)
487488
if query_tags:
488489
# Per-statement query tags. The kernel serialises the
489490
# dict (None value -> bare key) into the SEA
@@ -590,9 +591,7 @@ def execute_command(
590591
# native exception) — wrap the construction so callers see a
591592
# mapped PEP 249 exception.
592593
try:
593-
return self._make_result_set(
594-
executed, cursor, command_id, row_limit=row_limit
595-
)
594+
return self._make_result_set(executed, cursor, command_id)
596595
except Exception as exc:
597596
raise _wrap_kernel_exception("execute_command", exc) from exc
598597

@@ -764,9 +763,7 @@ def get_execution_result(
764763
# ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
765764
# can raise — map that to PEP 249 too.
766765
try:
767-
return self._make_result_set(
768-
stream, cursor, command_id, row_limit=cursor.row_limit
769-
)
766+
return self._make_result_set(stream, cursor, command_id)
770767
except Exception as exc:
771768
raise _wrap_kernel_exception("get_execution_result", exc) from exc
772769

@@ -777,7 +774,6 @@ def _make_result_set(
777774
kernel_handle: Any,
778775
cursor: "Cursor",
779776
command_id: CommandId,
780-
row_limit: Optional[int] = None,
781777
) -> "ResultSet":
782778
"""Build a ``KernelResultSet`` from any kernel handle. Used
783779
by sync execute, ``get_execution_result``, and all metadata
@@ -799,7 +795,6 @@ def _make_result_set(
799795
command_id=command_id,
800796
arraysize=cursor.arraysize,
801797
buffer_size_bytes=cursor.buffer_size_bytes,
802-
row_limit=row_limit,
803798
)
804799

805800
def _synthetic_command_id(self) -> CommandId:

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

Lines changed: 33 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,6 @@
2121
within a batch when ``n`` is smaller than the kernel's natural
2222
batch size; ``fetchall`` drains the whole stream.
2323
24-
When a cursor has ``row_limit`` set, this class caps the logical stream
25-
before rows reach any of the row or Arrow fetch APIs.
26-
2724
Note: ``buffer_size_bytes`` is accepted by the constructor for
2825
contract compatibility with the base ``ResultSet`` but is not
2926
consulted — the kernel backend currently caps buffering by rows
@@ -70,7 +67,6 @@ def __init__(
7067
command_id: CommandId,
7168
arraysize: int,
7269
buffer_size_bytes: int,
73-
row_limit: Optional[int] = None,
7470
):
7571
try:
7672
schema = kernel_handle.arrow_schema()
@@ -104,55 +100,27 @@ def __init__(
104100
# stays O(1) instead of walking the deque.
105101
self._buffered_count: int = 0
106102
self._exhausted: bool = False
107-
# The PyO3 kernel surface does not currently expose the core
108-
# StatementSpec row_limit setter. Enforce the cursor contract at
109-
# this streaming boundary until it does. Negative values retain the
110-
# existing unlimited behaviour; zero is a real zero-row limit.
111-
self._row_limit: Optional[int] = (
112-
row_limit if row_limit is not None and row_limit >= 0 else None
113-
)
114-
if self._row_limit == 0:
115-
self._mark_exhausted()
116103

117104
# ----- internal helpers -----
118105

119-
def _mark_exhausted(self) -> None:
120-
self._exhausted = True
121-
self.has_more_rows = False
122-
self.status = CommandState.SUCCEEDED
123-
124-
def _remaining_row_limit(self) -> Optional[int]:
125-
if self._row_limit is None:
126-
return None
127-
return max(
128-
0,
129-
self._row_limit - self._next_row_index - self._buffered_count,
130-
)
131-
132106
def _pull_one_batch(self) -> bool:
133107
"""Pull the next batch from the kernel into the local buffer.
134108
Returns True if a batch was added; False if the kernel side
135109
is exhausted."""
136110
if self._exhausted:
137111
return False
138-
remaining_limit = self._remaining_row_limit()
139-
if remaining_limit == 0:
140-
self._mark_exhausted()
141-
return False
142112
try:
143113
batch = self._kernel_handle.fetch_next_batch()
144114
except Exception as exc:
145115
raise wrap_kernel_exception("fetch_next_batch", exc) from exc
146116
if batch is None:
147-
self._mark_exhausted()
117+
self._exhausted = True
118+
self.has_more_rows = False
119+
self.status = CommandState.SUCCEEDED
148120
return False
149-
if remaining_limit is not None and batch.num_rows > remaining_limit:
150-
batch = batch.slice(0, remaining_limit)
151121
if batch.num_rows > 0:
152122
self._buffer.append(batch)
153123
self._buffered_count += batch.num_rows
154-
if remaining_limit is not None and batch.num_rows >= remaining_limit:
155-
self._mark_exhausted()
156124
return True
157125

158126
def _ensure_buffered(self, n_rows: int) -> int:
@@ -188,10 +156,36 @@ def _take_buffered(self, n: int) -> pyarrow.Table:
188156
return pyarrow.Table.from_batches(slices, schema=self._schema)
189157

190158
def _drain(self) -> pyarrow.Table:
191-
"""Consume the remaining logical stream into one table."""
192-
while not self._exhausted:
193-
self._pull_one_batch()
194-
return self._take_buffered(self._buffered_count)
159+
"""Consume everything left in the buffer + kernel stream
160+
and return as a single Table."""
161+
chunks: List[pyarrow.RecordBatch] = []
162+
if self._buffer and self._buffer_offset > 0:
163+
head = self._buffer.popleft()
164+
chunks.append(
165+
head.slice(self._buffer_offset, head.num_rows - self._buffer_offset)
166+
)
167+
self._buffer_offset = 0
168+
while self._buffer:
169+
chunks.append(self._buffer.popleft())
170+
if not self._exhausted:
171+
while True:
172+
try:
173+
batch = self._kernel_handle.fetch_next_batch()
174+
except Exception as exc:
175+
raise wrap_kernel_exception("fetch_next_batch", exc) from exc
176+
if batch is None:
177+
self._exhausted = True
178+
self.has_more_rows = False
179+
self.status = CommandState.SUCCEEDED
180+
break
181+
if batch.num_rows > 0:
182+
chunks.append(batch)
183+
rows = sum(c.num_rows for c in chunks)
184+
self._buffered_count = 0
185+
self._next_row_index += rows
186+
if not chunks:
187+
return pyarrow.Table.from_batches([], schema=self._schema)
188+
return pyarrow.Table.from_batches(chunks, schema=self._schema)
195189

196190
# ----- Arrow fetches -----
197191

‎tests/e2e/test_kernel_backend.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -183,11 +183,12 @@ def test_fetchall_arrow(conn):
183183
assert table.column_names == ["a", "b"]
184184

185185

186-
@pytest.mark.parametrize("row_limit", [0, 1, 5])
186+
@pytest.mark.parametrize("row_limit", [None, 0, 1, 5])
187187
def test_cursor_row_limit(conn, row_limit):
188188
with conn.cursor(row_limit=row_limit) as cur:
189189
cur.execute("SELECT id FROM range(10) ORDER BY id")
190-
assert [row[0] for row in cur.fetchall()] == list(range(row_limit))
190+
expected = list(range(10 if row_limit is None else row_limit))
191+
assert [row[0] for row in cur.fetchall()] == expected
191192

192193

193194
# ─── Logging (Rust kernel -> Python logging bridge) ──────────────────────────

‎tests/unit/test_kernel_client.py‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -436,7 +436,8 @@ def test_execute_command_forwards_query_tags():
436436
assert stmt.execute.called
437437

438438

439-
def test_execute_command_applies_row_limit_to_result_set():
439+
@pytest.mark.parametrize("row_limit", [None, 0, 1, 5])
440+
def test_execute_command_forwards_row_limit(row_limit):
440441
c = _make_client()
441442
c._kernel_session = MagicMock()
442443
cursor = MagicMock()
@@ -450,7 +451,7 @@ def test_execute_command_applies_row_limit_to_result_set():
450451
)
451452
c._kernel_session.statement.return_value = stmt
452453

453-
result = c.execute_command(
454+
c.execute_command(
454455
operation="SELECT * FROM range(10)",
455456
session_id=MagicMock(),
456457
max_rows=1,
@@ -461,11 +462,10 @@ def test_execute_command_applies_row_limit_to_result_set():
461462
parameters=[],
462463
async_op=False,
463464
enforce_embedded_schema_correctness=False,
464-
row_limit=5,
465+
row_limit=row_limit,
465466
)
466467

467-
assert result is not None
468-
assert result._row_limit == 5
468+
stmt.set_row_limit.assert_called_once_with(row_limit)
469469

470470

471471
# ---------------------------------------------------------------------------
@@ -809,13 +809,11 @@ def test_get_execution_result_attaches_by_id():
809809
cursor = MagicMock()
810810
cursor.arraysize = 100
811811
cursor.buffer_size_bytes = 1024
812-
cursor.row_limit = 5
813812
cid = CommandId.from_sea_statement_id("async-1")
814813

815814
rs = c.get_execution_result(cid, cursor=cursor)
816815

817816
assert rs is not None
818-
assert rs._row_limit == 5
819817
c._kernel_session.attach_async_statement.assert_called_with("async-1")
820818
handle.await_result.assert_called_once_with()
821819

@@ -1067,7 +1065,6 @@ def test_get_execution_result_is_re_callable():
10671065
cursor = MagicMock()
10681066
cursor.arraysize = 100
10691067
cursor.buffer_size_bytes = 1024
1070-
cursor.row_limit = None
10711068

10721069
rs1 = c.get_execution_result(cid, cursor=cursor)
10731070
rs2 = c.get_execution_result(cid, cursor=cursor)

‎tests/unit/test_kernel_result_set.py‎

Lines changed: 1 addition & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -28,15 +28,13 @@ def __init__(self, schema: pa.Schema, batches):
2828
self._schema = schema
2929
self._batches: Deque[pa.RecordBatch] = deque(batches)
3030
self.closed = False
31-
self.fetch_calls = 0
3231

3332
def arrow_schema(self) -> pa.Schema:
3433
return self._schema
3534

3635
def fetch_next_batch(self):
3736
if self.closed:
3837
raise RuntimeError("fetched after close")
39-
self.fetch_calls += 1
4038
if not self._batches:
4139
return None
4240
return self._batches.popleft()
@@ -45,7 +43,7 @@ def close(self):
4543
self.closed = True
4644

4745

48-
def _make_rs(handle, row_limit=None) -> KernelResultSet:
46+
def _make_rs(handle) -> KernelResultSet:
4947
# The base ResultSet __init__ takes a `connection` ref it never
5048
# actually dereferences during these buffer tests, so a Mock is
5149
# fine.
@@ -58,7 +56,6 @@ def _make_rs(handle, row_limit=None) -> KernelResultSet:
5856
command_id=CommandId.from_sea_statement_id("smoke-test"),
5957
arraysize=100,
6058
buffer_size_bytes=1024,
61-
row_limit=row_limit,
6259
)
6360

6461

@@ -143,63 +140,6 @@ def test_fetchall_rows(int_schema):
143140
assert [r[0] for r in rows] == [1, 2, 3]
144141

145142

146-
@pytest.mark.parametrize("row_limit", [0, 1, 5])
147-
def test_row_limit_caps_fetchall(int_schema, row_limit):
148-
handle = _FakeKernelHandle(
149-
int_schema,
150-
[_batch(int_schema, [0, 1, 2]), _batch(int_schema, list(range(3, 10)))],
151-
)
152-
rs = _make_rs(handle, row_limit=row_limit)
153-
154-
rows = rs.fetchall()
155-
156-
assert [row[0] for row in rows] == list(range(row_limit))
157-
assert rs.rownumber == row_limit
158-
159-
160-
def test_row_limit_applies_across_fetch_methods(int_schema):
161-
handle = _FakeKernelHandle(
162-
int_schema,
163-
[_batch(int_schema, [0, 1, 2]), _batch(int_schema, [3, 4, 5, 6])],
164-
)
165-
rs = _make_rs(handle, row_limit=5)
166-
167-
first = rs.fetchmany(2)
168-
third = rs.fetchone()
169-
rest = rs.fetchall_arrow()
170-
171-
assert [row[0] for row in first] == [0, 1]
172-
assert third is not None and third[0] == 2
173-
assert rest.column(0).to_pylist() == [3, 4]
174-
assert rs.fetchone() is None
175-
176-
177-
def test_row_limit_stops_before_fetching_extra_batches(int_schema):
178-
handle = _FakeKernelHandle(
179-
int_schema,
180-
[_batch(int_schema, [0, 1, 2]), _batch(int_schema, [3, 4, 5])],
181-
)
182-
rs = _make_rs(handle, row_limit=2)
183-
184-
assert rs.fetchall_arrow().column(0).to_pylist() == [0, 1]
185-
assert handle.fetch_calls == 1
186-
187-
188-
def test_row_limit_exact_batch_boundary_skips_exhaustion_fetch(int_schema):
189-
handle = _FakeKernelHandle(int_schema, [_batch(int_schema, [0, 1, 2])])
190-
rs = _make_rs(handle, row_limit=3)
191-
192-
assert rs.fetchall_arrow().column(0).to_pylist() == [0, 1, 2]
193-
assert handle.fetch_calls == 1
194-
195-
196-
def test_row_limit_larger_than_result_returns_all_rows(int_schema):
197-
handle = _FakeKernelHandle(int_schema, [_batch(int_schema, [1, 2, 3])])
198-
rs = _make_rs(handle, row_limit=10)
199-
200-
assert rs.fetchall_arrow().column(0).to_pylist() == [1, 2, 3]
201-
202-
203143
def test_fetchmany_negative_raises(int_schema):
204144
rs = _make_rs(_FakeKernelHandle(int_schema, []))
205145
with pytest.raises(ValueError):

0 commit comments

Comments
 (0)