|
1 | 1 | import unittest |
2 | 2 | import pytest |
3 | | -from unittest.mock import Mock |
| 3 | +from unittest.mock import Mock, patch |
4 | 4 |
|
5 | 5 | try: |
6 | 6 | import pyarrow as pa |
|
9 | 9 |
|
10 | 10 | import databricks.sql.client as client |
11 | 11 | from databricks.sql.backend.types import ExecuteResponse |
12 | | -from databricks.sql.utils import ArrowQueue |
| 12 | +from databricks.sql.utils import ArrowQueue, ColumnQueue, ColumnTable |
13 | 13 | from databricks.sql.backend.thrift_backend import ThriftDatabricksClient |
14 | 14 | from databricks.sql.result_set import ThriftResultSet |
15 | 15 |
|
@@ -391,5 +391,117 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self): |
391 | 391 | self.assertEqual(result.num_rows, 0) |
392 | 392 |
|
393 | 393 |
|
| 394 | +class FetchRowsPerRequestTests(unittest.TestCase): |
| 395 | + """maxRows sent by ThriftResultSet._fill_results_buffer for inline results.""" |
| 396 | + |
| 397 | + @staticmethod |
| 398 | + def _make_result_set(n_rows, arraysize, columnar=False): |
| 399 | + """A result set over ids 0..n_rows-1 from a fake server that returns exactly |
| 400 | + maxRows rows per FetchResults. Returns the result set and the maxRows sent.""" |
| 401 | + calls = [] |
| 402 | + next_row = 0 |
| 403 | + |
| 404 | + def fetch_results(**kwargs): |
| 405 | + nonlocal next_row |
| 406 | + calls.append(kwargs["max_rows"]) |
| 407 | + values = list(range(next_row, min(n_rows, next_row + kwargs["max_rows"]))) |
| 408 | + next_row += len(values) |
| 409 | + if columnar: |
| 410 | + queue = ColumnQueue(ColumnTable([values], ["col0"])) |
| 411 | + else: |
| 412 | + queue = ArrowQueue(pa.table({"col0": values}), len(values)) |
| 413 | + return queue, next_row < n_rows, 0 |
| 414 | + |
| 415 | + backend = Mock(spec=ThriftDatabricksClient) |
| 416 | + backend.fetch_results = fetch_results |
| 417 | + rs = ThriftResultSet( |
| 418 | + connection=Mock(), |
| 419 | + execute_response=ExecuteResponse( |
| 420 | + command_id=None, |
| 421 | + status=None, |
| 422 | + has_been_closed_server_side=False, |
| 423 | + description=[("col0", "integer", None, None, None, None, None)], |
| 424 | + lz4_compressed=True, |
| 425 | + is_staging_operation=False, |
| 426 | + ), |
| 427 | + thrift_client=backend, |
| 428 | + arraysize=arraysize, |
| 429 | + ) |
| 430 | + return rs, calls |
| 431 | + |
| 432 | + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") |
| 433 | + def test_fetchone_and_iteration_keep_arraysize(self): |
| 434 | + rs, calls = self._make_result_set(25, arraysize=10) |
| 435 | + |
| 436 | + rows = list(rs) |
| 437 | + |
| 438 | + self.assertEqual([r[0] for r in rows], list(range(25))) |
| 439 | + self.assertEqual(calls, [10, 10, 10]) |
| 440 | + |
| 441 | + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") |
| 442 | + def test_fetchmany_up_to_arraysize_keeps_arraysize(self): |
| 443 | + rs, calls = self._make_result_set(25, arraysize=10) |
| 444 | + |
| 445 | + ids = [r[0] for _ in range(4) for r in rs.fetchmany(7)] |
| 446 | + |
| 447 | + self.assertEqual(ids, list(range(25))) |
| 448 | + self.assertEqual(calls, [10, 10, 10]) |
| 449 | + |
| 450 | + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") |
| 451 | + def test_fetchall_requests_up_to_the_cap(self): |
| 452 | + rs, calls = self._make_result_set(25, arraysize=10) |
| 453 | + |
| 454 | + table = rs.fetchall_arrow() |
| 455 | + |
| 456 | + self.assertEqual(table.column(0).to_pylist(), list(range(25))) |
| 457 | + self.assertEqual(calls, [10, 2_000_000]) |
| 458 | + |
| 459 | + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") |
| 460 | + def test_large_fetchmany_requests_the_remaining_rows(self): |
| 461 | + rs, calls = self._make_result_set(100, arraysize=10) |
| 462 | + |
| 463 | + first = rs.fetchmany_arrow(35) |
| 464 | + second = rs.fetchmany_arrow(35) |
| 465 | + |
| 466 | + self.assertEqual(first.column(0).to_pylist(), list(range(35))) |
| 467 | + self.assertEqual(second.column(0).to_pylist(), list(range(35, 70))) |
| 468 | + self.assertEqual(calls, [10, 25, 35]) |
| 469 | + |
| 470 | + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") |
| 471 | + def test_rows_per_request_are_capped(self): |
| 472 | + with patch("databricks.sql.result_set._MAX_ROWS_PER_FETCH", 50): |
| 473 | + rs, calls = self._make_result_set(200, arraysize=10) |
| 474 | + table = rs.fetchmany_arrow(120) |
| 475 | + |
| 476 | + self.assertEqual(table.column(0).to_pylist(), list(range(120))) |
| 477 | + self.assertEqual(calls, [10, 50, 50, 10]) |
| 478 | + |
| 479 | + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") |
| 480 | + def test_arraysize_above_the_cap_is_kept(self): |
| 481 | + with patch("databricks.sql.result_set._MAX_ROWS_PER_FETCH", 50): |
| 482 | + rs, calls = self._make_result_set(200, arraysize=80) |
| 483 | + table = rs.fetchall_arrow() |
| 484 | + |
| 485 | + self.assertEqual(table.num_rows, 200) |
| 486 | + self.assertEqual(calls, [80, 80, 80]) |
| 487 | + |
| 488 | + def test_columnar_fetchall_requests_up_to_the_cap(self): |
| 489 | + rs, calls = self._make_result_set(25, arraysize=10, columnar=True) |
| 490 | + |
| 491 | + rows = rs.fetchall() |
| 492 | + |
| 493 | + self.assertEqual([r[0] for r in rows], list(range(25))) |
| 494 | + self.assertEqual(calls, [10, 2_000_000]) |
| 495 | + |
| 496 | + def test_columnar_fetchone_and_fetchmany(self): |
| 497 | + rs, calls = self._make_result_set(40, arraysize=10, columnar=True) |
| 498 | + |
| 499 | + ids = [rs.fetchone()[0]] + [r[0] for r in rs.fetchmany(30)] |
| 500 | + ids += [r[0] for r in rs.fetchmany(5)] |
| 501 | + |
| 502 | + self.assertEqual(ids, list(range(36))) |
| 503 | + self.assertEqual(calls, [10, 21, 10]) |
| 504 | + |
| 505 | + |
394 | 506 | if __name__ == "__main__": |
395 | 507 | unittest.main() |
0 commit comments