Skip to content

Commit ceb025c

Browse files
committed
Revert "UNPICK changes to review"
This reverts commit 2c800f3.
1 parent 2c800f3 commit ceb025c

File tree

18 files changed

+534
-165
lines changed

18 files changed

+534
-165
lines changed

docs/source/contributor-guide/ffi.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ as performant as possible and to utilize the features of DataFusion, you may dec
3434
your source in Rust and then expose it through `PyO3 <https://pyo3.rs>`_ as a Python library.
3535

3636
At first glance, it may appear the best way to do this is to add the ``datafusion-python``
37-
crate as a dependency, provide a ``PyTable``, and then to register it with the
37+
crate as a dependency, produce a DataFusion table in Rust, and then register it with the
3838
``SessionContext``. Unfortunately, this will not work.
3939

4040
When you produce your code as a Python library and it needs to interact with the DataFusion

docs/source/user-guide/data-sources.rst

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,11 @@ as Delta Lake. This will require a recent version of
152152
.. code-block:: python
153153
154154
from deltalake import DeltaTable
155+
from datafusion import TableProvider
155156
156157
delta_table = DeltaTable("path_to_table")
157-
ctx.register_table_provider("my_delta_table", delta_table)
158+
provider = TableProvider.from_capsule(delta_table.__datafusion_table_provider__())
159+
ctx.register_table("my_delta_table", provider)
158160
df = ctx.table("my_delta_table")
159161
df.show()
160162

docs/source/user-guide/io/table_provider.rst

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,20 +39,40 @@ A complete example can be found in the `examples folder <https://github.com/apac
3939
) -> PyResult<Bound<'py, PyCapsule>> {
4040
let name = CString::new("datafusion_table_provider").unwrap();
4141
42-
let provider = Arc::new(self.clone())
43-
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
44-
let provider = FFI_TableProvider::new(Arc::new(provider), false);
42+
let provider = Arc::new(self.clone());
43+
let provider = FFI_TableProvider::new(provider, false, None);
4544
4645
PyCapsule::new_bound(py, provider, Some(name.clone()))
4746
}
4847
}
4948
50-
Once you have this library available, in python you can register your table provider
51-
to the ``SessionContext``.
49+
Once you have this library available, you can construct a
50+
:py:class:`~datafusion.TableProvider` in Python and register it with the
51+
``SessionContext``. Table providers can be created either from the PyCapsule exposed by
52+
your Rust provider or from an existing :py:class:`~datafusion.dataframe.DataFrame`.
53+
Call the provider's ``__datafusion_table_provider__()`` method to obtain the capsule
54+
before constructing a ``TableProvider``. The ``TableProvider.from_view()`` helper is
55+
deprecated; instead use ``TableProvider.from_dataframe()`` or ``DataFrame.into_view()``.
5256

5357
.. code-block:: python
5458
59+
from datafusion import SessionContext, TableProvider
60+
61+
ctx = SessionContext()
5562
provider = MyTableProvider()
56-
ctx.register_table_provider("my_table", provider)
5763
58-
ctx.table("my_table").show()
64+
capsule = provider.__datafusion_table_provider__()
65+
capsule_provider = TableProvider.from_capsule(capsule)
66+
67+
df = ctx.from_pydict({"a": [1]})
68+
view_provider = TableProvider.from_dataframe(df)
69+
# or: view_provider = df.into_view()
70+
71+
ctx.register_table("capsule_table", capsule_provider)
72+
ctx.register_table("view_table", view_provider)
73+
74+
ctx.table("capsule_table").show()
75+
ctx.table("view_table").show()
76+
77+
Both ``TableProvider.from_capsule()`` and ``TableProvider.from_dataframe()`` create
78+
table providers that can be registered with the SessionContext using ``register_table()``.

examples/datafusion-ffi-example/python/tests/_test_table_function.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ def test_ffi_table_function_call_directly():
5353
table_udtf = udtf(table_func, "my_table_func")
5454

5555
my_table = table_udtf()
56-
ctx.register_table_provider("t", my_table)
56+
ctx.register_table("t", my_table)
5757
result = ctx.table("t").collect()
5858

5959
assert len(result) == 2

examples/datafusion-ffi-example/python/tests/_test_table_provider.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,16 @@
1818
from __future__ import annotations
1919

2020
import pyarrow as pa
21-
from datafusion import SessionContext
21+
from datafusion import SessionContext, TableProvider
2222
from datafusion_ffi_example import MyTableProvider
2323

2424

2525
def test_table_loading():
2626
ctx = SessionContext()
2727
table = MyTableProvider(3, 2, 4)
28-
ctx.register_table_provider("t", table)
28+
ctx.register_table(
29+
"t", TableProvider.from_capsule(table.__datafusion_table_provider__())
30+
)
2931
result = ctx.table("t").collect()
3032

3133
assert len(result) == 4

python/datafusion/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
from .io import read_avro, read_csv, read_json, read_parquet
5555
from .plan import ExecutionPlan, LogicalPlan
5656
from .record_batch import RecordBatch, RecordBatchStream
57+
from .table_provider import TableProvider
5758
from .user_defined import (
5859
Accumulator,
5960
AggregateUDF,
@@ -90,6 +91,7 @@
9091
"SessionContext",
9192
"Table",
9293
"TableFunction",
94+
"TableProvider",
9395
"WindowFrame",
9496
"WindowUDF",
9597
"catalog",

python/datafusion/catalog.py

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@
2727
if TYPE_CHECKING:
2828
import pyarrow as pa
2929

30+
from datafusion import TableProvider
31+
from datafusion.context import TableProviderExportable
32+
3033
try:
3134
from warnings import deprecated # Python 3.13+
3235
except ImportError:
@@ -122,8 +125,14 @@ def table(self, name: str) -> Table:
122125
"""Return the table with the given ``name`` from this schema."""
123126
return Table(self._raw_schema.table(name))
124127

125-
def register_table(self, name, table) -> None:
126-
"""Register a table provider in this schema."""
128+
def register_table(
129+
self, name, table: Table | TableProvider | TableProviderExportable
130+
) -> None:
131+
"""Register a table or table provider in this schema.
132+
133+
Objects implementing ``__datafusion_table_provider__`` are also supported
134+
and treated as :class:`TableProvider` instances.
135+
"""
127136
if isinstance(table, Table):
128137
return self._raw_schema.register_table(name, table.table)
129138
return self._raw_schema.register_table(name, table)
@@ -219,11 +228,16 @@ def table(self, name: str) -> Table | None:
219228
"""Retrieve a specific table from this schema."""
220229
...
221230

222-
def register_table(self, name: str, table: Table) -> None: # noqa: B027
223-
"""Add a table from this schema.
231+
def register_table( # noqa: B027
232+
self, name: str, table: Table | TableProvider | TableProviderExportable
233+
) -> None:
234+
"""Add a table to this schema.
224235
225236
This method is optional. If your schema provides a fixed list of tables, you do
226237
not need to implement this method.
238+
239+
Objects implementing ``__datafusion_table_provider__`` are also supported
240+
and treated as :class:`TableProvider` instances.
227241
"""
228242

229243
def deregister_table(self, name, cascade: bool) -> None: # noqa: B027

python/datafusion/context.py

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import pandas as pd
4747
import polars as pl
4848

49+
from datafusion import TableProvider
4950
from datafusion.plan import ExecutionPlan, LogicalPlan
5051

5152

@@ -734,7 +735,7 @@ def from_polars(self, data: pl.DataFrame, name: str | None = None) -> DataFrame:
734735
# https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116
735736
# is the discussion on how we arrived at adding register_view
736737
def register_view(self, name: str, df: DataFrame) -> None:
737-
"""Register a :py:class: `~datafusion.detaframe.DataFrame` as a view.
738+
"""Register a :py:class:`~datafusion.dataframe.DataFrame` as a view.
738739
739740
Args:
740741
name (str): The name to register the view under.
@@ -743,16 +744,31 @@ def register_view(self, name: str, df: DataFrame) -> None:
743744
view = df.into_view()
744745
self.ctx.register_table(name, view)
745746

746-
def register_table(self, name: str, table: Table) -> None:
747-
"""Register a :py:class: `~datafusion.catalog.Table` as a table.
747+
def register_table(
748+
self, name: str, table: Table | TableProvider | TableProviderExportable
749+
) -> None:
750+
"""Register a :py:class:`~datafusion.catalog.Table` or ``TableProvider``.
748751
749-
The registered table can be referenced from SQL statement executed against.
752+
The registered table can be referenced from SQL statements executed against
753+
this context.
754+
755+
Plain :py:class:`~datafusion.dataframe.DataFrame` objects are not supported;
756+
convert them first with :meth:`datafusion.dataframe.DataFrame.into_view` or
757+
:meth:`datafusion.catalog.TableProvider.from_dataframe`.
758+
759+
Objects implementing ``__datafusion_table_provider__`` are also supported
760+
and treated as :class:`~datafusion.catalog.TableProvider` instances.
750761
751762
Args:
752763
name: Name of the resultant table.
753-
table: DataFusion table to add to the session context.
764+
table: DataFusion :class:`Table`, :class:`TableProvider`, or any object
765+
implementing ``__datafusion_table_provider__`` to add to the session
766+
context.
754767
"""
755-
self.ctx.register_table(name, table.table)
768+
if isinstance(table, Table):
769+
self.ctx.register_table(name, table.table)
770+
else:
771+
self.ctx.register_table(name, table)
756772

757773
def deregister_table(self, name: str) -> None:
758774
"""Remove a table from the session."""
@@ -772,14 +788,21 @@ def register_catalog_provider(
772788
self.ctx.register_catalog_provider(name, provider)
773789

774790
def register_table_provider(
775-
self, name: str, provider: TableProviderExportable
791+
self, name: str, provider: Table | TableProvider | TableProviderExportable
776792
) -> None:
777793
"""Register a table provider.
778794
779-
This table provider must have a method called ``__datafusion_table_provider__``
780-
which returns a PyCapsule that exposes a ``FFI_TableProvider``.
795+
Deprecated: use :meth:`register_table` instead.
796+
797+
Objects implementing ``__datafusion_table_provider__`` are also supported
798+
and treated as :class:`~datafusion.catalog.TableProvider` instances.
781799
"""
782-
self.ctx.register_table_provider(name, provider)
800+
warnings.warn(
801+
"register_table_provider is deprecated; use register_table",
802+
DeprecationWarning,
803+
stacklevel=2,
804+
)
805+
self.register_table(name, provider)
783806

784807
def register_udtf(self, func: TableFunction) -> None:
785808
"""Register a user defined table function."""

python/datafusion/dataframe.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
import pyarrow as pa
5454

5555
from datafusion._internal import expr as expr_internal
56+
from datafusion.table_provider import TableProvider
5657

5758
from enum import Enum
5859

@@ -307,9 +308,17 @@ def __init__(self, df: DataFrameInternal) -> None:
307308
"""
308309
self.df = df
309310

310-
def into_view(self) -> pa.Table:
311-
"""Convert DataFrame as a ViewTable which can be used in register_table."""
312-
return self.df.into_view()
311+
def into_view(self) -> TableProvider:
312+
"""Convert ``DataFrame`` into a ``TableProvider`` view for registration.
313+
314+
This is the preferred way to obtain a view for
315+
:py:meth:`~datafusion.context.SessionContext.register_table`.
316+
``TableProvider.from_dataframe`` calls this method under the hood,
317+
and the older ``TableProvider.from_view`` helper is deprecated.
318+
"""
319+
from datafusion.table_provider import TableProvider as _TableProvider
320+
321+
return _TableProvider(self.df.into_view())
313322

314323
def __getitem__(self, key: str | list[str]) -> DataFrame:
315324
"""Return a new :py:class`DataFrame` with the specified column or columns.
Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""Wrapper helpers for :mod:`datafusion._internal.TableProvider`."""
18+
19+
from __future__ import annotations
20+
21+
from typing import Any
22+
23+
import datafusion._internal as df_internal
24+
25+
_InternalTableProvider = df_internal.TableProvider
26+
27+
28+
class TableProvider:
29+
"""High level wrapper around :mod:`datafusion._internal.TableProvider`."""
30+
31+
__slots__ = ("_table_provider",)
32+
33+
def __init__(self, table_provider: _InternalTableProvider) -> None:
34+
"""Wrap a low level :class:`~datafusion._internal.TableProvider`."""
35+
if isinstance(table_provider, TableProvider):
36+
table_provider = table_provider._table_provider
37+
38+
if not isinstance(table_provider, _InternalTableProvider):
39+
msg = "Expected a datafusion._internal.TableProvider instance."
40+
raise TypeError(msg)
41+
42+
self._table_provider = table_provider
43+
44+
# ------------------------------------------------------------------
45+
# constructors
46+
# ------------------------------------------------------------------
47+
@classmethod
48+
def _wrap(cls, provider: _InternalTableProvider | TableProvider) -> TableProvider:
49+
if isinstance(provider, cls):
50+
return provider
51+
return cls(provider)
52+
53+
@classmethod
54+
def from_capsule(cls, capsule: Any) -> TableProvider:
55+
"""Create a :class:`TableProvider` from a PyCapsule."""
56+
provider = _InternalTableProvider.from_capsule(capsule)
57+
return cls(provider)
58+
59+
@classmethod
60+
def from_dataframe(cls, df: Any) -> TableProvider:
61+
"""Create a :class:`TableProvider` from a :class:`DataFrame`."""
62+
from datafusion.dataframe import DataFrame as DataFrameWrapper
63+
64+
if isinstance(df, DataFrameWrapper):
65+
df = df.df
66+
67+
provider = _InternalTableProvider.from_dataframe(df)
68+
return cls(provider)
69+
70+
@classmethod
71+
def from_view(cls, df: Any) -> TableProvider:
72+
"""Create a :class:`TableProvider` from a DataFrame view."""
73+
from datafusion.dataframe import DataFrame as DataFrameWrapper
74+
75+
if isinstance(df, DataFrameWrapper):
76+
df = df.df
77+
78+
provider = _InternalTableProvider.from_view(df)
79+
return cls(provider)
80+
81+
# ------------------------------------------------------------------
82+
# passthrough helpers
83+
# ------------------------------------------------------------------
84+
def __getattr__(self, name: str) -> Any:
85+
"""Delegate attribute lookup to the wrapped provider."""
86+
return getattr(self._table_provider, name)
87+
88+
def __repr__(self) -> str: # pragma: no cover - simple delegation
89+
"""Return a representation of the wrapped provider."""
90+
return repr(self._table_provider)
91+
92+
def __datafusion_table_provider__(self) -> Any:
93+
"""Expose the wrapped provider for FFI integrations."""
94+
return self._table_provider.__datafusion_table_provider__()
95+
96+
97+
__all__ = ["TableProvider"]

0 commit comments

Comments
 (0)