11import abc
22import datetime
33import logging
4+ import os
45import threading
56import typing
67from collections import deque
2930_DUCKDB_BUFFER_THRESHOLD_CURRENT = _DUCKDB_BUFFER_THRESHOLD_ORIGINAL
3031_USE_DUCKDB_STATE = True
3132
33+ # DuckDB sizes its worker-thread pool to the host core count by default, spawning (threads - 1)
34+ # background workers per database instance. To keep the worker count from scaling with the number of
35+ # Struct channels (and with the host core count), every DuckDBState shares a single process-wide DuckDB
36+ # instance -- each state operates on its own cursor -- so the pool is created once and its size is a
37+ # single global knob. The default caps the pool at 4 threads: bounded on large hosts, while still
38+ # allowing some intra-query parallelism. Tune it with modify_duckdb_threads.
39+ # See https://github.com/Point72/csp-gateway/issues/292
40+ _DUCKDB_THREADS_ORIGINAL = min (os .cpu_count () or 1 , 4 )
41+ _DUCKDB_THREADS_CURRENT = _DUCKDB_THREADS_ORIGINAL
42+ # The shared DuckDB database instance, lazily created on first use. All DuckDBState objects operate on
43+ # cursors of this instance so that they share one TaskScheduler / worker-thread pool.
44+ _DUCKDB_SHARED_CONNECTION = None
45+ _DUCKDB_SHARED_CONNECTION_LOCK = threading .Lock ()
46+
3247log = logging .getLogger (__name__ )
3348
3449__all__ = [
3954 "build_track_state_node" ,
4055 "modify_buffer_threshold" ,
4156 "restore_buffer_threshold" ,
57+ "modify_duckdb_threads" ,
58+ "restore_duckdb_threads" ,
4259 "enable_duckdb_state" ,
4360 "disable_duckdb_state" ,
4461]
@@ -61,6 +78,49 @@ def restore_buffer_threshold() -> None:
6178 _DUCKDB_BUFFER_THRESHOLD_CURRENT = _DUCKDB_BUFFER_THRESHOLD_ORIGINAL
6279
6380
81+ def _get_duckdb_cursor () -> "duckdb.DuckDBPyConnection" :
82+ """Return a cursor on the process-wide shared DuckDB instance.
83+
84+ A single database instance is shared across all DuckDBState objects so that DuckDB's TaskScheduler
85+ (and therefore its worker-thread pool) is created once for the whole process rather than once per
86+ Struct channel. Each state uses its own cursor, which has an independent result context and so can
87+ be used safely from that state's own threads -- sharing a single connection object instead raises
88+ "Attempting to execute an unsuccessful or closed pending query result" under concurrent access.
89+ """
90+
91+ global _DUCKDB_SHARED_CONNECTION
92+ with _DUCKDB_SHARED_CONNECTION_LOCK :
93+ if _DUCKDB_SHARED_CONNECTION is None :
94+ _DUCKDB_SHARED_CONNECTION = duckdb .connect (config = {"threads" : _DUCKDB_THREADS_CURRENT })
95+ # The memory filesystem (used for bulk loading) is registered on the shared instance and is
96+ # visible to every cursor, so it must be registered exactly once here rather than per state.
97+ _DUCKDB_SHARED_CONNECTION .register_filesystem (fsspec .filesystem ("memory" ))
98+ return _DUCKDB_SHARED_CONNECTION .cursor ()
99+
100+
101+ def modify_duckdb_threads (new_threads : int ) -> None :
102+ """Set the size of the shared DuckDB worker-thread pool.
103+
104+ DuckDB spawns ``new_threads - 1`` background worker threads for the whole process (all state
105+ tables share one pool), so this is a single global knob independent of the host core count and the
106+ number of state channels. Takes effect immediately, including on an already-created shared instance.
107+ """
108+
109+ global _DUCKDB_THREADS_CURRENT
110+ if new_threads < 1 :
111+ raise ValueError ("Number of threads should be a positive integer" )
112+ with _DUCKDB_SHARED_CONNECTION_LOCK :
113+ _DUCKDB_THREADS_CURRENT = new_threads
114+ if _DUCKDB_SHARED_CONNECTION is not None :
115+ _DUCKDB_SHARED_CONNECTION .execute (f"SET threads={ new_threads } )
116+
117+
118+ def restore_duckdb_threads () -> None :
119+ """Reset the shared DuckDB worker-thread pool size to the original value"""
120+
121+ modify_duckdb_threads (_DUCKDB_THREADS_ORIGINAL )
122+
123+
64124def enable_duckdb_state () -> None :
65125 """Enable using DuckDBState as State class specialization"""
66126
@@ -205,10 +265,13 @@ def __init__(self, typ: Struct, keyby: Union[Tuple[str, ...], str], schema: Dict
205265 # Store the schema
206266 self ._schema = schema
207267 self ._schema_str = orjson .dumps (self ._schema ).decode ()
208- # NOTE: We make a new connection object for each state object. Using the same state connection across
209- # multiple state objects leads to issues in DuckDB's python API related to pending query results not
210- # being fully executed. This should be okay since we are doing this once during the initialization
211- self ._con = duckdb .connect ()
268+ # NOTE: Each state operates on its own cursor of a single process-wide DuckDB instance (see
269+ # _get_duckdb_cursor). We use a per-state cursor -- rather than one shared connection object --
270+ # because DuckDB's Python API raises on pending/unfinished query results when a single
271+ # connection is used concurrently from multiple states. Sharing the underlying instance keeps
272+ # the worker-thread pool bounded globally instead of one pool per Struct channel. The table
273+ # name (self._table_name) is unique per instance, so states never collide in the shared catalog.
274+ self ._con = _get_duckdb_cursor ()
212275 self ._keyby = keyby
213276 # ID generator for new records. We cannot rely on the id in the record as that might or might not exist
214277 self ._obj_id_generator = Counter (1 )
@@ -260,9 +323,6 @@ def __init__(self, typ: Struct, keyby: Union[Tuple[str, ...], str], schema: Dict
260323 # Open the memory file for buffering data before bulk loading
261324 self ._mem_file = fsspec .filesystem ("memory" ).open (self ._mem_file_name , "wb" )
262325
263- # Register memory filesystem so that duckdb can read from memory
264- self ._con .register_filesystem (fsspec .filesystem ("memory" ))
265-
266326 # TODO: Improve this to support querying dicts and lists as well, currently it only
267327 # works with structs and primitive types
268328 @classmethod
0 commit comments