Skip to content

Commit a6d6ee5

Browse files
ai: apply changes for #1007 (2 review threads)
Addresses: - #4419479174 at src/databricks/sql/backend/sea/utils/filters.py:179 - #4419479176 at src/databricks/sql/backend/kernel/result_set.py:211 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent 313aa5f commit a6d6ee5

2 files changed

Lines changed: 127 additions & 3 deletions

File tree

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

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -209,9 +209,15 @@ def _fetch_available(self, size: int) -> pyarrow.Table:
209209
return self._take_buffered(min(size, self._buffered_rows()))
210210

211211
@staticmethod
212-
def _prepend_buffered(
212+
def _concat_buffered_arrow(
213213
buffered: Optional[pyarrow.Table], rest: pyarrow.Table
214214
) -> pyarrow.Table:
215+
"""Prepend rows drained from the fetchone() buffer to ``rest``.
216+
217+
Both tables are built against the result set's ``_schema``, so a plain
218+
``pyarrow.concat_tables`` suffices (no ColumnTable / schema
219+
promotion handling as in ``ThriftResultSet._prepend_buffered``).
220+
"""
215221
if buffered is None:
216222
return rest
217223
if rest.num_rows == 0:
@@ -220,7 +226,7 @@ def _prepend_buffered(
220226

221227
def fetchall_arrow(self) -> pyarrow.Table:
222228
buffered = self._row_buffer_take_raw()
223-
return self._prepend_buffered(buffered, self._drain())
229+
return self._concat_buffered_arrow(buffered, self._drain())
224230

225231
def fetchmany_arrow(self, size: int) -> pyarrow.Table:
226232
if size < 0:
@@ -229,7 +235,7 @@ def fetchmany_arrow(self, size: int) -> pyarrow.Table:
229235
remaining = size - (buffered.num_rows if buffered is not None else 0)
230236
if buffered is not None and remaining == 0:
231237
return buffered
232-
return self._prepend_buffered(
238+
return self._concat_buffered_arrow(
233239
buffered, self._fetchmany_arrow_unbuffered(remaining)
234240
)
235241

‎tests/unit/test_filters.py‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
import unittest
66
from unittest.mock import MagicMock, patch
77

8+
import pytest
9+
810
from databricks.sql.backend.sea.utils.filters import ResultSetFilter
911

1012

@@ -153,5 +155,121 @@ def test_filter_tables_by_type(self):
153155
)
154156

155157

158+
class TestResultSetFilterWithFetchoneBuffer(unittest.TestCase):
159+
"""Rows buffered by fetchone() on a real SeaResultSet must reach the filter."""
160+
161+
N_ROWS = 12
162+
163+
def _row_type(self, i):
164+
return "TABLE" if i % 2 == 0 else "VIEW"
165+
166+
def _make_result_set(self, queue, result_format):
167+
from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
168+
from databricks.sql.backend.sea.result_set import SeaResultSet
169+
from databricks.sql.backend.types import CommandId, CommandState
170+
171+
connection = MagicMock()
172+
connection.open = True
173+
connection.disable_pandas = True
174+
sea_client = MagicMock()
175+
sea_client.max_download_threads = 10
176+
177+
execute_response = MagicMock()
178+
execute_response.command_id = CommandId.from_sea_statement_id("stmt-1")
179+
execute_response.status = CommandState.SUCCEEDED
180+
execute_response.has_been_closed_server_side = False
181+
execute_response.description = [
182+
("table_name", "string", None, None, None, None, None),
183+
("table_type", "string", None, None, None, None, None),
184+
]
185+
execute_response.is_staging_operation = False
186+
execute_response.lz4_compressed = False
187+
execute_response.arrow_schema_bytes = None
188+
189+
manifest = ResultManifest(
190+
format=result_format.value,
191+
schema={},
192+
total_row_count=self.N_ROWS,
193+
total_byte_count=-1,
194+
total_chunk_count=-1,
195+
)
196+
with patch(
197+
"databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue",
198+
return_value=queue,
199+
):
200+
return SeaResultSet(
201+
connection=connection,
202+
execute_response=execute_response,
203+
sea_client=sea_client,
204+
result_data=ResultData(data=None, external_links=None),
205+
manifest=manifest,
206+
buffer_size_bytes=1000,
207+
arraysize=100,
208+
)
209+
210+
def _consume_two_rows(self, rs):
211+
# 1st fetchone() refills 1 row, 2nd refills 10: rows 2..10 stay buffered
212+
# and row 11 is still in the queue.
213+
self.assertEqual(rs.fetchone()[0], "t0")
214+
self.assertEqual(rs.fetchone()[0], "t1")
215+
self.assertGreater(rs._row_buffer_count(), 0)
216+
217+
def _expected_names(self):
218+
return [
219+
"t{}".format(i)
220+
for i in range(2, self.N_ROWS)
221+
if self._row_type(i) == "TABLE"
222+
]
223+
224+
def test_json_filter_includes_rows_buffered_by_fetchone(self):
225+
from databricks.sql.backend.sea.queue import JsonQueue
226+
from databricks.sql.backend.sea.utils.constants import ResultFormat
227+
228+
data = [["t{}".format(i), self._row_type(i)] for i in range(self.N_ROWS)]
229+
rs = self._make_result_set(JsonQueue(data), ResultFormat.JSON_ARRAY)
230+
self._consume_two_rows(rs)
231+
232+
with patch(
233+
"databricks.sql.backend.sea.result_set.SeaResultSet"
234+
) as mock_sea_result_set_class:
235+
ResultSetFilter._filter_json_result_set(
236+
rs, 1, ["TABLE"], case_sensitive=True
237+
)
238+
result_data = mock_sea_result_set_class.call_args.kwargs["result_data"]
239+
240+
self.assertEqual([r[0] for r in result_data.data], self._expected_names())
241+
self.assertEqual(rs.rownumber, self.N_ROWS)
242+
243+
def test_arrow_filter_includes_rows_buffered_by_fetchone(self):
244+
pyarrow = pytest.importorskip("pyarrow")
245+
from databricks.sql.backend.sea.utils.constants import ResultFormat
246+
from databricks.sql.utils import ArrowQueue
247+
248+
table = pyarrow.Table.from_pydict(
249+
{
250+
"table_name": ["t{}".format(i) for i in range(self.N_ROWS)],
251+
"table_type": [self._row_type(i) for i in range(self.N_ROWS)],
252+
}
253+
)
254+
rs = self._make_result_set(
255+
ArrowQueue(table, self.N_ROWS), ResultFormat.ARROW_STREAM
256+
)
257+
self._consume_two_rows(rs)
258+
259+
with patch(
260+
"databricks.sql.backend.sea.result_set.SeaResultSet"
261+
) as mock_sea_result_set_class:
262+
ResultSetFilter._filter_arrow_result_set(
263+
rs, 1, ["TABLE"], case_sensitive=True
264+
)
265+
result_data = mock_sea_result_set_class.call_args.kwargs["result_data"]
266+
267+
filtered = pyarrow.ipc.open_stream(result_data.attachment).read_all()
268+
self.assertEqual(
269+
filtered.column("table_name").to_pylist(), self._expected_names()
270+
)
271+
self.assertEqual(rs.rownumber, self.N_ROWS)
272+
273+
156274
if __name__ == "__main__":
157275
unittest.main()

0 commit comments

Comments
 (0)