Skip to content

Commit 6011e65

Browse files
committed
Switch from RuntimeConfig to RuntimeEnvBuilder
1 parent 05802ab commit 6011e65

File tree

8 files changed

+52
-39
lines changed

8 files changed

+52
-39
lines changed

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ It is possible to configure runtime (memory and disk settings) and configuration
8787

8888
```python
8989
runtime = (
90-
RuntimeConfig()
90+
RuntimeEnvBuilder()
9191
.with_disk_manager_os()
9292
.with_fair_spill_pool(10000000)
9393
)

benchmarks/db-benchmark/groupby-datafusion.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
from datafusion import (
2323
col,
2424
functions as f,
25-
RuntimeConfig,
25+
RuntimeEnvBuilder,
2626
SessionConfig,
2727
SessionContext,
2828
)
@@ -85,7 +85,9 @@ def execute(df):
8585

8686
# create a session context with explicit runtime and config settings
8787
runtime = (
88-
RuntimeConfig().with_disk_manager_os().with_fair_spill_pool(64 * 1024 * 1024 * 1024)
88+
RuntimeEnvBuilder()
89+
.with_disk_manager_os()
90+
.with_fair_spill_pool(64 * 1024 * 1024 * 1024)
8991
)
9092
config = (
9193
SessionConfig()

benchmarks/tpch/tpch.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def bench(data_path, query_path):
2828

2929
# create context
3030
# runtime = (
31-
# RuntimeConfig()
31+
# RuntimeEnvBuilder()
3232
# .with_disk_manager_os()
3333
# .with_fair_spill_pool(10000000)
3434
# )

docs/source/user-guide/configuration.rst

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,18 @@ Configuration
1919
=============
2020

2121
Let's look at how we can configure DataFusion. When creating a :py:class:`~datafusion.context.SessionContext`, you can pass in
22-
a :py:class:`~datafusion.context.SessionConfig` and :py:class:`~datafusion.context.RuntimeConfig` object. These two cover a wide range of options.
22+
a :py:class:`~datafusion.context.SessionConfig` and :py:class:`~datafusion.context.RuntimeEnvBuilder` object. These two cover a wide range of options.
2323

2424
.. code-block:: python
2525
26-
from datafusion import RuntimeConfig, SessionConfig, SessionContext
26+
from datafusion import RuntimeEnvBuilder, SessionConfig, SessionContext
2727
2828
# create a session context with default settings
2929
ctx = SessionContext()
3030
print(ctx)
3131
3232
# create a session context with explicit runtime and config settings
33-
runtime = RuntimeConfig().with_disk_manager_os().with_fair_spill_pool(10000000)
33+
runtime = RuntimeEnvBuilder().with_disk_manager_os().with_fair_spill_pool(10000000)
3434
config = (
3535
SessionConfig()
3636
.with_create_default_catalog_and_schema(True)
@@ -48,4 +48,4 @@ a :py:class:`~datafusion.context.SessionConfig` and :py:class:`~datafusion.conte
4848
4949
5050
You can read more about available :py:class:`~datafusion.context.SessionConfig` options in the `rust DataFusion Configuration guide <https://arrow.apache.org/datafusion/user-guide/configs.html>`_,
51-
and about :code:`RuntimeConfig` options in the rust `online API documentation <https://docs.rs/datafusion/latest/datafusion/execution/runtime_env/struct.RuntimeConfig.html>`_.
51+
and about :code:`RuntimeEnvBuilder` options in the rust `online API documentation <https://docs.rs/datafusion/latest/datafusion/execution/runtime_env/struct.RuntimeEnvBuilder.html>`_.

examples/create-context.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,14 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
from datafusion import RuntimeConfig, SessionConfig, SessionContext
18+
from datafusion import RuntimeEnvBuilder, SessionConfig, SessionContext
1919

2020
# create a session context with default settings
2121
ctx = SessionContext()
2222
print(ctx)
2323

2424
# create a session context with explicit runtime and config settings
25-
runtime = RuntimeConfig().with_disk_manager_os().with_fair_spill_pool(10000000)
25+
runtime = RuntimeEnvBuilder().with_disk_manager_os().with_fair_spill_pool(10000000)
2626
config = (
2727
SessionConfig()
2828
.with_create_default_catalog_and_schema(True)

python/datafusion/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
from .context import (
3030
SessionContext,
3131
SessionConfig,
32-
RuntimeConfig,
32+
RuntimeEnvBuilder,
3333
SQLOptions,
3434
)
3535

@@ -66,7 +66,7 @@
6666
"SessionContext",
6767
"SessionConfig",
6868
"SQLOptions",
69-
"RuntimeConfig",
69+
"RuntimeEnvBuilder",
7070
"Expr",
7171
"ScalarUDF",
7272
"WindowFrame",

python/datafusion/context.py

Lines changed: 33 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from __future__ import annotations
2121

2222
from ._internal import SessionConfig as SessionConfigInternal
23-
from ._internal import RuntimeConfig as RuntimeConfigInternal
23+
from ._internal import RuntimeEnvBuilder as RuntimeEnvBuilderInternal
2424
from ._internal import SQLOptions as SQLOptionsInternal
2525
from ._internal import SessionContext as SessionContextInternal
2626

@@ -256,56 +256,58 @@ def set(self, key: str, value: str) -> SessionConfig:
256256
return self
257257

258258

259-
class RuntimeConfig:
259+
class RuntimeEnvBuilder:
260260
"""Runtime configuration options."""
261261

262262
def __init__(self) -> None:
263-
"""Create a new :py:class:`RuntimeConfig` with default values."""
264-
self.config_internal = RuntimeConfigInternal()
263+
"""Create a new :py:class:`RuntimeEnvBuilder` with default values."""
264+
self.config_internal = RuntimeEnvBuilderInternal()
265265

266-
def with_disk_manager_disabled(self) -> RuntimeConfig:
266+
def with_disk_manager_disabled(self) -> RuntimeEnvBuilder:
267267
"""Disable the disk manager, attempts to create temporary files will error.
268268
269269
Returns:
270-
A new :py:class:`RuntimeConfig` object with the updated setting.
270+
A new :py:class:`RuntimeEnvBuilder` object with the updated setting.
271271
"""
272272
self.config_internal = self.config_internal.with_disk_manager_disabled()
273273
return self
274274

275-
def with_disk_manager_os(self) -> RuntimeConfig:
275+
def with_disk_manager_os(self) -> RuntimeEnvBuilder:
276276
"""Use the operating system's temporary directory for disk manager.
277277
278278
Returns:
279-
A new :py:class:`RuntimeConfig` object with the updated setting.
279+
A new :py:class:`RuntimeEnvBuilder` object with the updated setting.
280280
"""
281281
self.config_internal = self.config_internal.with_disk_manager_os()
282282
return self
283283

284-
def with_disk_manager_specified(self, *paths: str | pathlib.Path) -> RuntimeConfig:
284+
def with_disk_manager_specified(
285+
self, *paths: str | pathlib.Path
286+
) -> RuntimeEnvBuilder:
285287
"""Use the specified paths for the disk manager's temporary files.
286288
287289
Args:
288290
paths: Paths to use for the disk manager's temporary files.
289291
290292
Returns:
291-
A new :py:class:`RuntimeConfig` object with the updated setting.
293+
A new :py:class:`RuntimeEnvBuilder` object with the updated setting.
292294
"""
293295
paths_list = [str(p) for p in paths]
294296
self.config_internal = self.config_internal.with_disk_manager_specified(
295297
paths_list
296298
)
297299
return self
298300

299-
def with_unbounded_memory_pool(self) -> RuntimeConfig:
301+
def with_unbounded_memory_pool(self) -> RuntimeEnvBuilder:
300302
"""Use an unbounded memory pool.
301303
302304
Returns:
303-
A new :py:class:`RuntimeConfig` object with the updated setting.
305+
A new :py:class:`RuntimeEnvBuilder` object with the updated setting.
304306
"""
305307
self.config_internal = self.config_internal.with_unbounded_memory_pool()
306308
return self
307309

308-
def with_fair_spill_pool(self, size: int) -> RuntimeConfig:
310+
def with_fair_spill_pool(self, size: int) -> RuntimeEnvBuilder:
309311
"""Use a fair spill pool with the specified size.
310312
311313
This pool works best when you know beforehand the query has multiple spillable
@@ -326,16 +328,16 @@ def with_fair_spill_pool(self, size: int) -> RuntimeConfig:
326328
size: Size of the memory pool in bytes.
327329
328330
Returns:
329-
A new :py:class:`RuntimeConfig` object with the updated setting.
331+
A new :py:class:`RuntimeEnvBuilder` object with the updated setting.
330332
331333
Examples usage::
332334
333-
config = RuntimeConfig().with_fair_spill_pool(1024)
335+
config = RuntimeEnvBuilder().with_fair_spill_pool(1024)
334336
"""
335337
self.config_internal = self.config_internal.with_fair_spill_pool(size)
336338
return self
337339

338-
def with_greedy_memory_pool(self, size: int) -> RuntimeConfig:
340+
def with_greedy_memory_pool(self, size: int) -> RuntimeEnvBuilder:
339341
"""Use a greedy memory pool with the specified size.
340342
341343
This pool works well for queries that do not need to spill or have a single
@@ -346,32 +348,39 @@ def with_greedy_memory_pool(self, size: int) -> RuntimeConfig:
346348
size: Size of the memory pool in bytes.
347349
348350
Returns:
349-
A new :py:class:`RuntimeConfig` object with the updated setting.
351+
A new :py:class:`RuntimeEnvBuilder` object with the updated setting.
350352
351353
Example usage::
352354
353-
config = RuntimeConfig().with_greedy_memory_pool(1024)
355+
config = RuntimeEnvBuilder().with_greedy_memory_pool(1024)
354356
"""
355357
self.config_internal = self.config_internal.with_greedy_memory_pool(size)
356358
return self
357359

358-
def with_temp_file_path(self, path: str | pathlib.Path) -> RuntimeConfig:
360+
def with_temp_file_path(self, path: str | pathlib.Path) -> RuntimeEnvBuilder:
359361
"""Use the specified path to create any needed temporary files.
360362
361363
Args:
362364
path: Path to use for temporary files.
363365
364366
Returns:
365-
A new :py:class:`RuntimeConfig` object with the updated setting.
367+
A new :py:class:`RuntimeEnvBuilder` object with the updated setting.
366368
367369
Example usage::
368370
369-
config = RuntimeConfig().with_temp_file_path("/tmp")
371+
config = RuntimeEnvBuilder().with_temp_file_path("/tmp")
370372
"""
371373
self.config_internal = self.config_internal.with_temp_file_path(str(path))
372374
return self
373375

374376

377+
@deprecated("Use `RuntimeEnvBuilder` instead.")
378+
class RuntimeConfig(RuntimeEnvBuilder):
379+
"""See `RuntimeEnvBuilder`."""
380+
381+
pass
382+
383+
375384
class SQLOptions:
376385
"""Options to be used when performing SQL queries."""
377386

@@ -445,7 +454,9 @@ class SessionContext:
445454
"""
446455

447456
def __init__(
448-
self, config: SessionConfig | None = None, runtime: RuntimeConfig | None = None
457+
self,
458+
config: SessionConfig | None = None,
459+
runtime: RuntimeEnvBuilder | None = None,
449460
) -> None:
450461
"""Main interface for executing queries with DataFusion.
451462

python/tests/test_context.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525

2626
from datafusion import (
2727
DataFrame,
28-
RuntimeConfig,
28+
RuntimeEnvBuilder,
2929
SessionConfig,
3030
SessionContext,
3131
SQLOptions,
@@ -43,7 +43,7 @@ def test_create_context_session_config_only():
4343

4444

4545
def test_create_context_runtime_config_only():
46-
SessionContext(runtime=RuntimeConfig())
46+
SessionContext(runtime=RuntimeEnvBuilder())
4747

4848

4949
@pytest.mark.parametrize("path_to_str", (True, False))
@@ -54,7 +54,7 @@ def test_runtime_configs(tmp_path, path_to_str):
5454
path1 = str(path1) if path_to_str else path1
5555
path2 = str(path2) if path_to_str else path2
5656

57-
runtime = RuntimeConfig().with_disk_manager_specified(path1, path2)
57+
runtime = RuntimeEnvBuilder().with_disk_manager_specified(path1, path2)
5858
config = SessionConfig().with_default_catalog_and_schema("foo", "bar")
5959
ctx = SessionContext(config, runtime)
6060
assert ctx is not None
@@ -67,7 +67,7 @@ def test_runtime_configs(tmp_path, path_to_str):
6767
def test_temporary_files(tmp_path, path_to_str):
6868
path = str(tmp_path) if path_to_str else tmp_path
6969

70-
runtime = RuntimeConfig().with_temp_file_path(path)
70+
runtime = RuntimeEnvBuilder().with_temp_file_path(path)
7171
config = SessionConfig().with_default_catalog_and_schema("foo", "bar")
7272
ctx = SessionContext(config, runtime)
7373
assert ctx is not None
@@ -77,7 +77,7 @@ def test_temporary_files(tmp_path, path_to_str):
7777

7878

7979
def test_create_context_with_all_valid_args():
80-
runtime = RuntimeConfig().with_disk_manager_os().with_fair_spill_pool(10000000)
80+
runtime = RuntimeEnvBuilder().with_disk_manager_os().with_fair_spill_pool(10000000)
8181
config = (
8282
SessionConfig()
8383
.with_create_default_catalog_and_schema(True)

0 commit comments

Comments
 (0)