From 62ae45176f985c9238b972b7d0a6389f435cfe21 Mon Sep 17 00:00:00 2001 From: Chongchen Chen Date: Sun, 22 Dec 2024 15:09:54 +0800 Subject: [PATCH 1/3] feat: support enable_url_table config --- examples/create-context.py | 4 ++++ python/datafusion/context.py | 19 +++++++++++++++++-- src/context.rs | 11 +++++++---- 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/examples/create-context.py b/examples/create-context.py index 3184d4085..898eb7972 100644 --- a/examples/create-context.py +++ b/examples/create-context.py @@ -37,3 +37,7 @@ ) ctx = SessionContext(config, runtime) print(ctx) + +config = config.with_url_table(True) +ctx = SessionContext(config, runtime) +print(ctx) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index a07b5d175..1e519a123 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -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 @@ -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 diff --git a/src/context.rs b/src/context.rs index 8675e97df..a32c7d82a 100644 --- a/src/context.rs +++ b/src/context.rs @@ -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, runtime: Option, + enable_url_table: bool, ) -> PyResult { let config = if let Some(c) = config { c.config @@ -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 From 99c2fd082503c1712a2eabfdeefc5ec8e92ad1e3 Mon Sep 17 00:00:00 2001 From: Chongchen Chen Date: Tue, 7 Jan 2025 22:36:18 +0800 Subject: [PATCH 2/3] change enable_url_table as method --- examples/create-context.py | 3 +-- python/datafusion/context.py | 31 ++++++++++++++----------------- src/context.rs | 17 ++++++++++------- 3 files changed, 25 insertions(+), 26 deletions(-) diff --git a/examples/create-context.py b/examples/create-context.py index 898eb7972..11525d8b8 100644 --- a/examples/create-context.py +++ b/examples/create-context.py @@ -38,6 +38,5 @@ ctx = SessionContext(config, runtime) print(ctx) -config = config.with_url_table(True) -ctx = SessionContext(config, runtime) +ctx = ctx.enable_url_table() print(ctx) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 1e519a123..02659adda 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -66,27 +66,13 @@ def __arrow_c_array__( # noqa: D105 class SessionConfig: """Session configuration options.""" - def __init__(self, config_options: dict[str, str] | None = None, enable_url_table: bool = False) -> None: + def __init__(self, config_options: dict[str, str] | None = None) -> 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 @@ -481,11 +467,22 @@ 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, enable_url_table) + self.ctx = SessionContextInternal(config, runtime) + + def enable_url_table(self) -> "SessionContext": + + """Control if local files can be queried as tables. + + Returns: + A new :py:class:`SessionContext` object with url table enabled. + """ + klass = self.__class__ + obj = klass.__new__(klass) + obj.ctx = self.ctx.enable_url_table() + return obj def register_object_store( self, schema: str, store: Any, host: str | None = None diff --git a/src/context.rs b/src/context.rs index a32c7d82a..9420331bc 100644 --- a/src/context.rs +++ b/src/context.rs @@ -272,12 +272,11 @@ pub struct PySessionContext { #[pymethods] impl PySessionContext { - #[pyo3(signature = (config=None, runtime=None, enable_url_table=false))] + #[pyo3(signature = (config=None, runtime=None))] #[new] pub fn new( config: Option, runtime: Option, - enable_url_table: bool, ) -> PyResult { let config = if let Some(c) = config { c.config @@ -295,11 +294,15 @@ impl PySessionContext { .with_runtime_env(runtime) .with_default_features() .build(); - let mut ctx = SessionContext::new_with_state(session_state); - if enable_url_table { - ctx = ctx.enable_url_table(); - } - Ok(PySessionContext { ctx }) + Ok(PySessionContext { + ctx: SessionContext::new_with_state(session_state), + }) + } + + pub fn enable_url_table(&self) -> PyResult { + Ok(PySessionContext { + ctx: self.ctx.clone().enable_url_table(), + }) } /// Register an object store with the given name From fe0472a62cacbf570c18431e081ebce195ca69e1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 8 Jan 2025 16:39:06 -0500 Subject: [PATCH 3/3] Remove whitespace --- python/datafusion/context.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 02659adda..e6ec261c0 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -473,7 +473,6 @@ def __init__( self.ctx = SessionContextInternal(config, runtime) def enable_url_table(self) -> "SessionContext": - """Control if local files can be queried as tables. Returns: