Skip to content

Commit

Permalink
feat: support enable_url_table config
Browse files Browse the repository at this point in the history
  • Loading branch information
chenkovsky committed Dec 22, 2024
1 parent 79c22d6 commit 62ae451
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 6 deletions.
4 changes: 4 additions & 0 deletions examples/create-context.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,3 +37,7 @@
)
ctx = SessionContext(config, runtime)
print(ctx)

config = config.with_url_table(True)
ctx = SessionContext(config, runtime)
print(ctx)
19 changes: 17 additions & 2 deletions python/datafusion/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,13 +66,27 @@ def __arrow_c_array__( # noqa: D105
class SessionConfig:
"""Session configuration options."""

def __init__(self, config_options: dict[str, str] | None = None) -> None:
def __init__(self, config_options: dict[str, str] | None = None, enable_url_table: bool = False) -> None:
"""Create a new :py:class:`SessionConfig` with the given configuration options.
Args:
config_options: Configuration options.
"""
self.config_internal = SessionConfigInternal(config_options)
self.enable_url_table = enable_url_table

def with_url_table(self, enabled: bool = True) -> SessionConfig:

"""Control if local files can be queried as tables.
Args:
enabled: Whether local files can be queried as tables.
Returns:
A new :py:class:`SessionConfig` object with the updated setting.
"""
self.enable_url_table = enabled
return self

def with_create_default_catalog_and_schema(
self, enabled: bool = True
Expand Down Expand Up @@ -467,10 +481,11 @@ def __init__(
ctx = SessionContext()
df = ctx.read_csv("data.csv")
"""
enable_url_table = config.enable_url_table if config is not None else False
config = config.config_internal if config is not None else None
runtime = runtime.config_internal if runtime is not None else None

self.ctx = SessionContextInternal(config, runtime)
self.ctx = SessionContextInternal(config, runtime, enable_url_table)

def register_object_store(
self, schema: str, store: Any, host: str | None = None
Expand Down
11 changes: 7 additions & 4 deletions src/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -272,11 +272,12 @@ pub struct PySessionContext {

#[pymethods]
impl PySessionContext {
#[pyo3(signature = (config=None, runtime=None))]
#[pyo3(signature = (config=None, runtime=None, enable_url_table=false))]
#[new]
pub fn new(
config: Option<PySessionConfig>,
runtime: Option<PyRuntimeConfig>,
enable_url_table: bool,
) -> PyResult<Self> {
let config = if let Some(c) = config {
c.config
Expand All @@ -294,9 +295,11 @@ impl PySessionContext {
.with_runtime_env(runtime)
.with_default_features()
.build();
Ok(PySessionContext {
ctx: SessionContext::new_with_state(session_state),
})
let mut ctx = SessionContext::new_with_state(session_state);
if enable_url_table {
ctx = ctx.enable_url_table();
}
Ok(PySessionContext { ctx })
}

/// Register an object store with the given name
Expand Down

0 comments on commit 62ae451

Please sign in to comment.