Add missing SessionContext utility methods by timsaucer · Pull Request #1475 · apache/datafusion-python

timsaucer · 2026-04-01T16:17:12Z

Which issue does this PR close?

Closes #1459

Rationale for this change

These methods exist in the upstream repository but have not been exposed to Python.

What changes are included in this PR?

Add methods to the Python API
Add unit tests

Are there any user-facing changes?

New addition only.

Expose upstream DataFusion v53 utility methods: session_start_time, enable_ident_normalization, parse_sql_expr, execute_logical_plan, refresh_catalogs, remove_optimizer_rule, and table_provider. The add_optimizer_rule and add_analyzer_rule methods are omitted as the OptimizerRule and AnalyzerRule traits are not yet exposed to Python. Closes apache#1459. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Copilot

Pull request overview

This PR exposes additional SessionContext utility/introspection methods in the datafusion-python API to match capabilities available in upstream DataFusion v53 (Issue #1459), and adds unit tests to cover the new Python surface area.

Changes:

Added Python SessionContext wrappers for: session_start_time, enable_ident_normalization, parse_sql_expr, execute_logical_plan, refresh_catalogs, remove_optimizer_rule, and table_provider.
Added corresponding Rust binding methods on PySessionContext to call into DataFusion v53 APIs.
Added unit tests validating the new Python methods.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
`python/datafusion/context.py`	Adds new `SessionContext` methods to the public Python API and wraps internal bindings (`Expr`, `DataFrame`, `Table`).
`crates/core/src/context.rs`	Exposes the underlying DataFusion `SessionContext` methods via the PyO3 `PySessionContext` bindings.
`python/tests/test_context.py`	Adds tests for the newly exposed `SessionContext` methods.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

crates/core/src/context.rs

python/datafusion/context.py

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

nuno-faria

Thanks @timsaucer.

nuno-faria · 2026-04-06T20:55:32Z

python/datafusion/context.py

+        Examples:
+            >>> ctx = SessionContext()
+            >>> start_time = ctx.session_start_time()
+            >>> assert "T" in start_time  # RFC 3339 contains a 'T' separator


This assert feels a little odd, what about showing a result?

>>> ctx.session_start_time() '2026-01-01T12:34:56.123456789+00:00'

nuno-faria · 2026-04-06T20:56:29Z

python/datafusion/context.py

+
+        Examples:
+            >>> ctx = SessionContext()
+            >>> assert isinstance(ctx.enable_ident_normalization(), bool)


Same thing here:

Suggested change

>>> assert isinstance(ctx.enable_ident_normalization(), bool)

>>> ctx.enable_ident_normalization()

True

nuno-faria · 2026-04-06T20:57:19Z

python/datafusion/context.py

+        from datafusion.expr import Expr  # noqa: PLC0415
+
+        return Expr(self.ctx.parse_sql_expr(sql, schema))


I think we could remove the import and the wrapping with Expr.

nuno-faria · 2026-04-06T20:57:53Z

python/datafusion/context.py

+        from datafusion.catalog import Table  # noqa: PLC0415
+
+        return Table(self.ctx.table_provider(name))


Also here I think we can remove the Table and the import.

nuno-faria · 2026-04-06T20:59:35Z

python/tests/test_context.py

+    assert isinstance(st, str)
+    assert "T" in st  # RFC 3339 format


What about this? The conversion should fail if the string is badly formatted.

Suggested change

assert isinstance(st, str)

assert "T" in st # RFC 3339 format

dt.datetime.fromisoformat(st).isoformat()

nuno-faria · 2026-04-06T21:01:03Z

python/datafusion/context.py

+            >>> expr = ctx.parse_sql_expr("1 + 2", schema)
+            >>> assert "Int64(1) + Int64(2)" in str(expr)


Suggested change

>>> expr = ctx.parse_sql_expr("1 + 2", schema)

>>> assert "Int64(1) + Int64(2)" in str(expr)

>>> ctx.parse_sql_expr("1 + 2", schema)

Expr(Int64(1) + Int64(2))

nuno-faria · 2026-04-06T21:01:23Z

python/tests/test_context.py

+
+    schema = DFSchema.empty()
+    expr = ctx.parse_sql_expr("1 + 2", schema)
+    assert "Int64(1) + Int64(2)" in str(expr)


Suggested change

assert "Int64(1) + Int64(2)" in str(expr)

assert str(expr) == "Expr(Int64(1) + Int64(2))"

nuno-faria · 2026-04-06T21:02:20Z

python/tests/test_context.py

+
+
+def test_remove_optimizer_rule(ctx):
+    assert ctx.remove_optimizer_rule("nonexistent_rule") is False


Testing with a rule that exists as well:

Suggested change

assert ctx.remove_optimizer_rule("nonexistent_rule") is False

assert ctx.remove_optimizer_rule("push_down_filter")

assert ctx.remove_optimizer_rule("nonexistent_rule") is False

nuno-faria · 2026-04-06T21:09:47Z

python/tests/test_context.py

+    result = ctx.enable_ident_normalization()
+    assert isinstance(result, bool)


I think it's better to change the value and check it.

Suggested change

result = ctx.enable_ident_normalization()

assert isinstance(result, bool)

assert ctx.enable_ident_normalization()

ctx.sql("set datafusion.sql_parser.enable_ident_normalization = false")

assert ctx.enable_ident_normalization() is False

Unrelated but the original method name is a bit misleading since it does not enable the flag, only returns the value.

timsaucer requested a review from Copilot April 3, 2026 19:17

Copilot started reviewing on behalf of timsaucer April 3, 2026 19:18 View session

Copilot AI reviewed Apr 3, 2026

View reviewed changes

crates/core/src/context.rs Outdated Show resolved Hide resolved

python/datafusion/context.py Show resolved Hide resolved

timsaucer and others added 2 commits April 5, 2026 08:14

Raise KeyError from table_provider for consistency with table()

f72830f

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Add docstring examples for new SessionContext utility methods

bb888f4

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

timsaucer marked this pull request as ready for review April 5, 2026 12:28

update docstring

c7f47f1

nuno-faria approved these changes Apr 6, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add missing SessionContext utility methods#1475

Add missing SessionContext utility methods#1475
timsaucer wants to merge 4 commits intoapache:mainfrom
timsaucer:feat/add-utility-methods

timsaucer commented Apr 1, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

nuno-faria left a comment

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

nuno-faria Apr 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	>>> assert isinstance(ctx.enable_ident_normalization(), bool)
	>>> ctx.enable_ident_normalization()
	True

		from datafusion.expr import Expr # noqa: PLC0415

		return Expr(self.ctx.parse_sql_expr(sql, schema))

		from datafusion.catalog import Table # noqa: PLC0415

		return Table(self.ctx.table_provider(name))

		assert isinstance(st, str)
		assert "T" in st # RFC 3339 format

	assert isinstance(st, str)
	assert "T" in st # RFC 3339 format
	dt.datetime.fromisoformat(st).isoformat()

		>>> expr = ctx.parse_sql_expr("1 + 2", schema)
		>>> assert "Int64(1) + Int64(2)" in str(expr)

	assert "Int64(1) + Int64(2)" in str(expr)
	assert str(expr) == "Expr(Int64(1) + Int64(2))"



		def test_remove_optimizer_rule(ctx):
		assert ctx.remove_optimizer_rule("nonexistent_rule") is False

	assert ctx.remove_optimizer_rule("nonexistent_rule") is False
	assert ctx.remove_optimizer_rule("push_down_filter")
	assert ctx.remove_optimizer_rule("nonexistent_rule") is False

		result = ctx.enable_ident_normalization()
		assert isinstance(result, bool)

-    result = ctx.enable_ident_normalization()
-    assert isinstance(result, bool)
+    assert ctx.enable_ident_normalization()
+    ctx.sql("set datafusion.sql_parser.enable_ident_normalization = false")
+    assert ctx.enable_ident_normalization() is False

Conversation

timsaucer commented Apr 1, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

nuno-faria left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants