Skip to content

Commit 62ae451

Browse files
committed
feat: support enable_url_table config
1 parent 79c22d6 commit 62ae451

File tree

3 files changed

+28
-6
lines changed

3 files changed

+28
-6
lines changed

examples/create-context.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,3 +37,7 @@
3737
)
3838
ctx = SessionContext(config, runtime)
3939
print(ctx)
40+
41+
config = config.with_url_table(True)
42+
ctx = SessionContext(config, runtime)
43+
print(ctx)

python/datafusion/context.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,13 +66,27 @@ def __arrow_c_array__( # noqa: D105
6666
class SessionConfig:
6767
"""Session configuration options."""
6868

69-
def __init__(self, config_options: dict[str, str] | None = None) -> None:
69+
def __init__(self, config_options: dict[str, str] | None = None, enable_url_table: bool = False) -> None:
7070
"""Create a new :py:class:`SessionConfig` with the given configuration options.
7171
7272
Args:
7373
config_options: Configuration options.
7474
"""
7575
self.config_internal = SessionConfigInternal(config_options)
76+
self.enable_url_table = enable_url_table
77+
78+
def with_url_table(self, enabled: bool = True) -> SessionConfig:
79+
80+
"""Control if local files can be queried as tables.
81+
82+
Args:
83+
enabled: Whether local files can be queried as tables.
84+
85+
Returns:
86+
A new :py:class:`SessionConfig` object with the updated setting.
87+
"""
88+
self.enable_url_table = enabled
89+
return self
7690

7791
def with_create_default_catalog_and_schema(
7892
self, enabled: bool = True
@@ -467,10 +481,11 @@ def __init__(
467481
ctx = SessionContext()
468482
df = ctx.read_csv("data.csv")
469483
"""
484+
enable_url_table = config.enable_url_table if config is not None else False
470485
config = config.config_internal if config is not None else None
471486
runtime = runtime.config_internal if runtime is not None else None
472487

473-
self.ctx = SessionContextInternal(config, runtime)
488+
self.ctx = SessionContextInternal(config, runtime, enable_url_table)
474489

475490
def register_object_store(
476491
self, schema: str, store: Any, host: str | None = None

src/context.rs

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -272,11 +272,12 @@ pub struct PySessionContext {
272272

273273
#[pymethods]
274274
impl PySessionContext {
275-
#[pyo3(signature = (config=None, runtime=None))]
275+
#[pyo3(signature = (config=None, runtime=None, enable_url_table=false))]
276276
#[new]
277277
pub fn new(
278278
config: Option<PySessionConfig>,
279279
runtime: Option<PyRuntimeConfig>,
280+
enable_url_table: bool,
280281
) -> PyResult<Self> {
281282
let config = if let Some(c) = config {
282283
c.config
@@ -294,9 +295,11 @@ impl PySessionContext {
294295
.with_runtime_env(runtime)
295296
.with_default_features()
296297
.build();
297-
Ok(PySessionContext {
298-
ctx: SessionContext::new_with_state(session_state),
299-
})
298+
let mut ctx = SessionContext::new_with_state(session_state);
299+
if enable_url_table {
300+
ctx = ctx.enable_url_table();
301+
}
302+
Ok(PySessionContext { ctx })
300303
}
301304

302305
/// Register an object store with the given name

0 commit comments

Comments
 (0)