Skip to content

Commit 99ba9d4

Browse files
committed
fix: enforce TypeError for DataFrame registration in SessionContext
1 parent 5de9b53 commit 99ba9d4

File tree

2 files changed

+10
-3
lines changed

2 files changed

+10
-3
lines changed

python/datafusion/context.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424

2525
import pyarrow as pa
2626

27+
from datafusion import EXPECTED_PROVIDER_MSG
28+
2729
try:
2830
from warnings import deprecated # Python 3.13+
2931
except ImportError:
@@ -767,8 +769,13 @@ def register_table(
767769
implementing ``__datafusion_table_provider__`` to add to the session
768770
context.
769771
"""
770-
provider = _normalize_table_provider(table)
771-
self.ctx.register_table(name, provider)
772+
if isinstance(table, DataFrame):
773+
raise TypeError(EXPECTED_PROVIDER_MSG)
774+
775+
if isinstance(table, Table):
776+
self.ctx.register_table(name, table.table)
777+
else:
778+
self.ctx.register_table(name, table)
772779

773780
def deregister_table(self, name: str) -> None:
774781
"""Remove a table from the session."""

python/tests/test_context.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -423,7 +423,7 @@ def test_table_provider_from_capsule_invalid():
423423

424424
def test_register_table_with_dataframe_errors(ctx):
425425
df = ctx.from_pydict({"a": [1]})
426-
with pytest.raises(Exception) as exc_info:
426+
with pytest.raises(TypeError) as exc_info:
427427
ctx.register_table("bad", df)
428428

429429
assert str(exc_info.value) == EXPECTED_PROVIDER_MSG

0 commit comments

Comments
 (0)