GeezerVison

The World Wide Web, restored to 800 pixels of glory.

GeezerVison is showing github.com — perf: pause the GC while rows are built (optional, stacked on #1039) by vuanhphung · Pull Request #1043 · databricks/databricks-sql-python

perf: pause the GC while rows are built (optional, stacked on #1039) #1043

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

vuanhphung wants to merge 3 commits into perf/fetchall-rows

base: perf/fetchall-rows

from perf/fetchall-rows-construct

+256 −4

Open

perf: pause the GC while rows are built (optional, stacked on #1039) #1043

Show file tree

Hide file tree

Changes from all commits

Commits

Show all changes

3 commits

Select commit Hold shift + click to select a range

056f35a

perf: pause the cyclic GC while each batch of rows is built

vuanhphung Nov 12, 2026

ef58abe

fix: undo a half-joined GC pause when __enter__ is interrupted

vuanhphung Nov 12, 2026

f1d4061

fix: clear the GC pause flag inside the lock; document the override

vuanhphung Nov 12, 2026

File filter

Filter by extension .md (1) .py (2) All 2 file types selected

Viewed files

Clear filters

Conversations

Failed to load comments.

Uh oh!

There was an error while loading. Please reload this page.

Jump to

Jump to file

Failed to load files.

Uh oh!

There was an error while loading. Please reload this page.

Diff view

Unified

Split

Hide whitespace

Diff view

Unified

Split

Hide whitespace

There are no files selected for viewing

1 change: 1 addition & 0 deletions 1 CHANGELOG.md

Show comments

View file

Open in desktop

Original file line number	Diff line number	Diff line change
Expand Up		@@ -4,6 +4,7 @@
		- Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629))
		- Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most.
		- Lower the peak memory of `fetchall()` and large `fetchmany()` calls by converting Arrow results to rows 100,000 rows at a time instead of building one DataFrame for the whole result, and build each row with one Python call instead of three, which makes these calls about 10% faster.
		- Speed up `fetchall()` and large `fetchmany()` calls on the default (pandas) conversion path by about another 30%: Python's cyclic garbage collector is paused while each batch of up to 100,000 rows is built (only if it is enabled; it is re-enabled as soon as no batch is being built, so between batches unless conversions overlap). If the application calls `gc.disable()` while a batch is being built, the GC is re-enabled when that pause ends; a GC that is already disabled when `fetchall()` runs is left alone.

		# 4.9.0 (2026-11-04)
		- Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes.
Expand Down

90 changes: 86 additions & 4 deletions 90 src/databricks/sql/result_set.py

Show comments

View file

Open in desktop

Original file line number	Diff line number	Diff line change
Expand Up		@@ -3,7 +3,12 @@
		from abc import ABC, abstractmethod
		from typing import List, Optional, TYPE_CHECKING, Tuple

		import gc
		import logging
		import os
		import threading
		import time

		import pandas

		try:
Expand Down Expand Up		@@ -55,6 +60,80 @@ def _converts_in_slices(arrow_type) -> bool:
		)


		# Each Row and its __dict__ are tracked by the cyclic GC, whose collections rescan
		# all rows built so far (about a third of the conversion time of a large result),
		# so rows are built with the GC paused, one bounded batch at a time.
		_GC_PAUSE_BATCH_ROWS = 100_000
		_GC_PAUSE_MIN_ROWS = 10_000
		_GC_PAUSE_MAX_SECONDS = 0.5


		class _GCPause:
		"""Pauses the cyclic GC while one batch of rows is built.

		Threads share the pause: the first batch disables the GC (only if it is
		enabled) and the last one re-enables it and collects the young generation, so
		other threads' cyclic garbage is still freed between batches. Batches stop
		joining a pause once it is _GC_PAUSE_MAX_SECONDS old; it ends when the batches
		already in it finish. Other code that disables the GC while a pause is running
		has it re-enabled when the pause ends: the gc module can't tell who disabled it.
		"""

		lock = threading.Lock()
		active = 0
		started = 0.0

		def __init__(self, num_rows: int):
		self._num_rows = num_rows
		self._joined = False

		def __enter__(self) -> None:
		if self._num_rows < _GC_PAUSE_MIN_ROWS:
		return
		try:
		with _GCPause.lock:
		now = time.monotonic()
		if _GCPause.active == 0:
		if not gc.isenabled():
		return
		_GCPause.started = now
		elif now - _GCPause.started > _GC_PAUSE_MAX_SECONDS:
		return
		# No call between these two lines, so no async exception can
		# split them.
		_GCPause.active += 1
		self._joined = True
		gc.disable()
Comment thread peco-review-bot[bot] marked this conversation as resolved. Resolved Uh oh! There was an error while loading. Please reload this page.
		except BaseException:
		# E.g. a KeyboardInterrupt after joining: __exit__ won't be called.
		self.__exit__(None, None, None)
		raise

		def __exit__(self, *exc_info) -> None:
		if not self._joined:
		return
		with _GCPause.lock:
		# No call between these two lines, so no async exception can split them.
		self._joined = False
		_GCPause.active -= 1
		if _GCPause.active:
		return
Comment thread peco-review-bot[bot] marked this conversation as resolved. Resolved Uh oh! There was an error while loading. Please reload this page. Comment thread peco-review-bot[bot] marked this conversation as resolved. Resolved Uh oh! There was an error while loading. Please reload this page.
		gc.enable()
		gc.collect(0)


		def _end_gc_pause_in_child() -> None:
		# Batches running in other threads of the parent never finish in the child.
		if _GCPause.active:
		_GCPause.active = 0
		gc.enable()
		_GCPause.lock = threading.Lock()


		if hasattr(os, "register_at_fork"):
		os.register_at_fork(after_in_child=_end_gc_pause_in_child)


		def _make_rows(ResultRow: Row, values) -> List[Row]:
		"""The rows ``[ResultRow(*v) for v in values]``, built with one Python call per
		row (``Row.__setattr__``) instead of three. ``values`` is a 2-D object array."""
Expand All		@@ -63,10 +142,13 @@ def _make_rows(ResultRow: Row, values) -> List[Row]:
		new_row = tuple.__new__
		rows: List[Row] = []
		append = rows.append
		for v in values:
		row = new_row(Row, v)
		row.__fields__ = ResultRow
		append(row)
		for start in range(0, len(values), _GC_PAUSE_BATCH_ROWS):
		batch = values[start : start + _GC_PAUSE_BATCH_ROWS]
		with _GCPause(len(values)):
		for v in batch:
		row = new_row(Row, v)
		row.__fields__ = ResultRow
		append(row)
		return rows


Expand Down

169 changes: 169 additions & 0 deletions 169 tests/unit/test_pandas_compatibility.py

Show comments

View file

Open in desktop

Original file line number	Diff line number	Diff line change
Expand Up		@@ -8,8 +8,12 @@

		import copy
		import datetime
		import gc
		import os
		import pickle
		import threading
		import unittest
		import warnings
		from decimal import Decimal
		from unittest.mock import Mock, patch

Expand Down Expand Up		@@ -643,5 +647,170 @@ def test_duration_next_to_nested_column(self):
		self.assertEqual(row[0], expected[0])


		@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
		class TestRowBuildingGCPause(unittest.TestCase):
		"""Rows are built with the cyclic GC paused per bounded batch."""

		def setUp(self):
		self.assertTrue(gc.isenabled())
		self.assertEqual(result_set._GCPause.active, 0)
		self.addCleanup(gc.enable)

		def _convert(self, num_rows=5, **patches):
		table = pa.table({"a": list(range(num_rows))})
		rs = _make_result_set([("a", "bigint", None, None, None, None, None)])
		patches = {"_GC_PAUSE_MIN_ROWS": 1, "_GC_PAUSE_BATCH_ROWS": 2, **patches}
		with patch.multiple(result_set, **patches), patch.object(
		result_set.gc, "disable", wraps=gc.disable
		) as disable, patch.object(
		result_set.gc, "collect", wraps=gc.collect
		) as collect:
		rows = rs._convert_arrow_table(table)
		self.assertEqual([r.a for r in rows], list(range(num_rows)))
		return disable.call_count, collect.call_count

		def test_each_batch_pauses_and_gc_is_reenabled(self):
		# 5 rows in batches of 2: three pauses, each ending with a young collection.
		self.assertEqual(self._convert(), (3, 3))
		self.assertTrue(gc.isenabled())

		def test_each_batch_of_each_slice_pauses(self):
		# Slices of 3 rows (3 + 2), built in batches of 2: 2 + 1 pauses.
		self.assertEqual(self._convert(_ROWS_PER_CONVERSION=3), (3, 3))

		def test_small_batches_do_not_pause(self):
		self.assertEqual(self._convert(_GC_PAUSE_MIN_ROWS=10_000), (0, 0))

		def test_gc_disabled_by_the_caller_stays_disabled(self):
		gc.disable()
		self.assertEqual(self._convert(), (0, 0))
		self.assertFalse(gc.isenabled())

		def test_interrupt_while_joining_a_pause_is_undone(self):
		real_disable = gc.disable

		def disable_then_interrupt():
		real_disable()
		raise KeyboardInterrupt

		with patch.object(result_set.gc, "disable", disable_then_interrupt):
		with self.assertRaises(KeyboardInterrupt):
		with result_set._GCPause(10**6):
		self.fail("the batch should not start")
		self.assertTrue(gc.isenabled())
		self.assertEqual(result_set._GCPause.active, 0)
		# Later pauses still work.
		with result_set._GCPause(10**6):
		self.assertFalse(gc.isenabled())
		self.assertTrue(gc.isenabled())

		def test_interrupt_while_ending_a_pause_keeps_state_consistent(self):
		real_enable = gc.enable

		def enable_then_interrupt():
		real_enable()
		raise KeyboardInterrupt

		pause = result_set._GCPause(10**6)
		pause.__enter__()
		with patch.object(result_set.gc, "enable", enable_then_interrupt):
		with self.assertRaises(KeyboardInterrupt):
		pause.__exit__(None, None, None)
		self.assertTrue(gc.isenabled())
		self.assertEqual(result_set._GCPause.active, 0)

		def test_gc_is_reenabled_when_building_fails(self):
		with self.assertRaises(ZeroDivisionError):
		with result_set._GCPause(10**6):
		self.assertFalse(gc.isenabled())
		1 / 0
		self.assertTrue(gc.isenabled())
		self.assertEqual(result_set._GCPause.active, 0)

		def test_overlapping_pauses_end_with_the_last_one(self):
		first, second = result_set._GCPause(106), result_set._GCPause(106)
		first.__enter__()
		second.__enter__()
		first.__exit__(None, None, None)
		self.assertFalse(gc.isenabled())
		second.__exit__(None, None, None)
		self.assertTrue(gc.isenabled())

		def test_pause_is_not_extended_past_the_limit(self):
		first, late = result_set._GCPause(106), result_set._GCPause(106)
		with patch.object(result_set, "_GC_PAUSE_MAX_SECONDS", -1):
		first.__enter__()
		late.__enter__() # the pause is already too old to be extended
		first.__exit__(None, None, None)
		self.assertTrue(gc.isenabled())
		late.__exit__(None, None, None)
		self.assertTrue(gc.isenabled())
		self.assertEqual(result_set._GCPause.active, 0)
Comment thread peco-review-bot[bot] marked this conversation as resolved. Resolved Uh oh! There was an error while loading. Please reload this page.

		def test_fork_child_ends_a_pause_left_by_other_threads(self):
		lock = result_set._GCPause.lock
		self.addCleanup(setattr, result_set._GCPause, "lock", lock)
		result_set._GCPause(10**6).__enter__()
		result_set._end_gc_pause_in_child()
		self.assertTrue(gc.isenabled())
		self.assertEqual(result_set._GCPause.active, 0)
		self.assertIsNot(result_set._GCPause.lock, lock)

		@unittest.skipUnless(hasattr(os, "fork"), "needs os.fork")
		def test_real_fork_child_starts_with_gc_enabled(self):
		# Dedicated exit codes: Python itself exits with 1 on an uncaught error.
		ok, gc_disabled, still_active, error = 70, 71, 72, 73
		pause = result_set._GCPause(10**6)
		pause.__enter__()
		try:
		with warnings.catch_warnings():
		# Python 3.12+ warns about forking a multi-threaded process.
		warnings.simplefilter("ignore", DeprecationWarning)
		pid = os.fork()
		if pid == 0:
		code = error
		try:
		if not gc.isenabled():
		code = gc_disabled
		elif result_set._GCPause.active:
		code = still_active
		else:
		code = ok
		finally:
		os._exit(code)
		finally:
		pause.__exit__(None, None, None)
		_, status = os.waitpid(pid, 0)
		self.assertEqual(os.waitstatus_to_exitcode(status), ok)
		self.assertTrue(gc.isenabled())

		def test_concurrent_conversions_leave_gc_enabled(self):
		table = pa.table({"a": list(range(50))})
		rs = _make_result_set([("a", "bigint", None, None, None, None, None)])
		errors = []

		def convert():
		try:
		for _ in range(20):
		assert len(rs._convert_arrow_table(table)) == 50
		except Exception as e: # pragma: no cover - reported below
		errors.append(e)

		with patch.multiple(
		result_set,
		_GC_PAUSE_MIN_ROWS=1,
		_GC_PAUSE_BATCH_ROWS=3,
		_ROWS_PER_CONVERSION=7,
		):
		threads = [threading.Thread(target=convert) for _ in range(4)]
		for t in threads:
		t.start()
		for t in threads:
		t.join()
		self.assertEqual(errors, [])
		self.assertTrue(gc.isenabled())
		self.assertEqual(result_set._GCPause.active, 0)


		if __name__ == "__main__":
		unittest.main()

Oops, something went wrong.

Uh oh!

There was an error while loading. Please reload this page.

Add this suggestion to a batch that can be applied as a single commit. This suggestion is invalid because no changes were made to the code. Suggestions cannot be applied while the pull request is closed. Suggestions cannot be applied while viewing a subset of changes. Only one suggestion per line can be applied in a batch. Add this suggestion to a batch that can be applied as a single commit. Applying suggestions on deleted lines is not supported. You must change the existing code in this line in order to create a valid suggestion. Outdated suggestions cannot be applied. This suggestion has been applied or marked resolved. Suggestions cannot be applied from pending reviews. Suggestions cannot be applied on multi-line comments. Suggestions cannot be applied while the pull request is queued to merge. Suggestion cannot be applied right now. Please check back later.

You are visitor 01402520 — GeezerVison does not store the pages you look up.
Best viewed with Netscape Navigator 4.7 or Internet Explorer 6 at 800×600. No cookies required.