batchcorder-datafusion 55.0.0__cp310-abi3-win_amd64.whl

batchcorder_datafusion/__init__.py

@@ -0,0 +1,36 @@
+"""
+DataFusion table provider that tees one read-once Arrow stream to every scan.
+
+Register a :class:`TeeTableProvider` with a ``SessionContext`` from a
+DataFusion host (datafusion-python or xorq-datafusion) and every scan of the table in a query replays the same
+upstream stream, which is read only once.  With the default
+``bounding="sealed"``, each planned scan reserves a reader at planning time,
+and each batch is evicted as soon as every scan has read past it.
+
+Example::
+
+    import pyarrow as pa
+    from datafusion import SessionContext
+    from batchcorder_datafusion import TeeTableProvider
+
+    ctx = SessionContext()
+    ctx.register_table("t", TeeTableProvider(reader))
+    ctx.sql("select * from t a join t b using (k)").collect()
+
+A sealed provider serves one query.  Register a new provider per query, or
+pass ``bounding="unbounded"`` for a table that is queried repeatedly.
+
+The host must be built on the same DataFusion major as this wheel.
+Registering the provider with any other host raises :class:`RuntimeError`.
+"""
+
+from __future__ import annotations
+
+from importlib.metadata import version
+
+from ._batchcorder_datafusion import TeeTableProvider
+
+
+__all__ = ["TeeTableProvider"]
+
+__version__: str = version("batchcorder-datafusion")

batchcorder_datafusion/_batchcorder_datafusion.pyd

@@ -0,0 +1,4 @@
+[diffend] Oversized file quarantined before diffing.
+name: batchcorder_datafusion/_batchcorder_datafusion.pyd
+size: 61756416 bytes
+sha256: f55e64eb0161e1b09d321b16bc203ca8b995beff41abb752d219112246627869

batchcorder_datafusion/_batchcorder_datafusion.pyi

@@ -0,0 +1,35 @@
+import typing
+
+import pyarrow as pa
+
+__all__ = ["TeeTableProvider"]
+
+@typing.final
+class TeeTableProvider:
+    def __new__(
+        cls,
+        reader: typing.Any,
+        *,
+        memory_capacity: int | None = None,
+        disk_path: str | None = None,
+        disk_capacity: int | None = None,
+        write_policy: str = "on_insertion",
+        bounding: str = "sealed",
+    ) -> TeeTableProvider: ...
+    def __datafusion_table_provider__(self, session: typing.Any) -> typing.Any: ...
+    @property
+    def schema(self) -> pa.Schema: ...
+    @property
+    def scans_planned(self) -> int: ...
+    @property
+    def scans_executed(self) -> int: ...
+    @property
+    def ingested_count(self) -> int: ...
+    @property
+    def evicted_count(self) -> int: ...
+    @property
+    def upstream_exhausted(self) -> bool: ...
+    @property
+    def is_sealed(self) -> bool: ...
+    def seal(self) -> None: ...
+    def close(self) -> None: ...

batchcorder_datafusion/_hosts.py

@@ -0,0 +1,168 @@
+"""Check that a DataFusion host is built on the same DataFusion major as this wheel."""
+
+from __future__ import annotations
+
+import re
+import sys
+from importlib.metadata import PackageNotFoundError, version
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+# The provider crosses into the host through datafusion-ffi, whose structs are
+# only compatible within one DataFusion major.  Each entry maps a host
+# distribution to the half-open range [low, high) of its versions built on the
+# same DataFusion major as this wheel (datafusion-ffi in Cargo.toml).
+#
+# Keep this table and the compatibility table in README.md in sync.
+SUPPORTED_HOSTS: dict[str, tuple[tuple[int, ...], tuple[int, ...]]] = {
+    "datafusion": ((55,), (56,)),
+    "xorq-datafusion": ((0, 3), (0, 4)),
+}
+
+# Top-level Python package of each host -> its distribution.  A session class
+# belongs to a host when its module is in the package, or when the package
+# exports it (see `_host_distribution`).
+HOST_PACKAGES: dict[str, str] = {
+    "datafusion": "datafusion",
+    "xorq_datafusion": "xorq-datafusion",
+}
+
+DATAFUSION_MAJOR = 55
+
+
+def _release(text: str) -> tuple[int, ...]:
+    """
+    Return the leading numeric release segment of a version string.
+
+    Parameters
+    ----------
+    text
+        A version such as ``"55.0.0"`` or ``"0.3.0rc1"``.
+
+    Returns
+    -------
+    tuple[int, ...]
+        The release numbers, e.g. ``(55, 0, 0)`` or ``(0, 3, 0)``.
+
+    """
+    match = re.match(r"\d+(?:\.\d+)*", text)
+    if match is None:
+        return ()
+    return tuple(int(part) for part in match.group().split("."))
+
+
+def _format(release: tuple[int, ...]) -> str:
+    return ".".join(map(str, release))
+
+
+def _exports(package: str, cls: type) -> bool:
+    """
+    Tell whether ``package`` or one of its imported submodules exports ``cls``.
+
+    Parameters
+    ----------
+    package
+        A host's top-level Python package.
+    cls
+        The class of the session the host passed.
+
+    Returns
+    -------
+    bool
+        ``True`` if an imported module of ``package`` has ``cls`` as an
+        attribute.
+
+    """
+    prefix = package + "."
+    # Only modules already imported: the host imported its own package to
+    # build the session, and importing one here could have side effects.
+    for name, module in list(sys.modules.items()):
+        if module is None or not (name == package or name.startswith(prefix)):
+            continue
+        if any(value is cls for value in vars(module).values()):
+            return True
+    return False
+
+
+def _host_distribution(session: object) -> str | None:
+    """
+    Return the distribution of the host that ``session`` belongs to.
+
+    Parameters
+    ----------
+    session
+        The ``session`` argument the host passed to
+        ``__datafusion_table_provider__``.
+
+    Returns
+    -------
+    str | None
+        The host's distribution, or ``None`` for an unknown host.
+
+    """
+    cls = type(session)
+    dist = HOST_PACKAGES.get(cls.__module__.partition(".")[0])
+    if dist is not None or cls.__module__ in sys.modules:
+        return dist
+    # A native (PyO3) class reports the module name given at build time,
+    # which need not be a real module: xorq-datafusion's SessionContext says
+    # `let`.  Find the host package that exports the class itself.  A class
+    # from a real module (a capsule is `builtins.PyCapsule`, which hosts
+    # re-export) is never looked up this way.
+    for package, dist in HOST_PACKAGES.items():
+        if _exports(package, cls):
+            return dist
+    return None
+
+
+def check_host(
+    session: object,
+    get_version: Callable[[str], str] = version,
+) -> None:
+    """
+    Raise if ``session`` belongs to a host built on another DataFusion major.
+
+    Called by ``TeeTableProvider.__datafusion_table_provider__`` before any
+    FFI struct of the host is touched.  The host is identified by the package
+    of the session's class, or by the host package that exports that class.
+    A codec capsule, or a session from an unknown host, cannot be identified
+    here; the provider then checks the datafusion-ffi version of the host's
+    codec instead.
+
+    Parameters
+    ----------
+    session
+        The ``session`` argument the host passed to
+        ``__datafusion_table_provider__``.
+    get_version
+        Callable mapping a distribution name to its installed version and
+        raising :class:`importlib.metadata.PackageNotFoundError` when it is
+        not installed.  Defaults to :func:`importlib.metadata.version`.
+
+    Raises
+    ------
+    RuntimeError
+        If the host's installed version is outside its supported range.
+
+    """
+    dist = _host_distribution(session)
+    if dist is None:
+        return
+    try:
+        installed = get_version(dist)
+    except PackageNotFoundError:
+        return
+    low, high = SUPPORTED_HOSTS[dist]
+    if low <= _release(installed) < high:
+        return
+    raise RuntimeError(
+        f"TeeTableProvider cannot be registered with {dist} {installed}: "
+        f"batchcorder-datafusion is built on DataFusion {DATAFUSION_MAJOR} and "
+        f"needs {dist}>={_format(low)},<{_format(high)}.  Install the "
+        f"batchcorder-datafusion release built on the same DataFusion major "
+        f"as {dist} {installed}, or a {dist} in that range."
+    )

batchcorder_datafusion-55.0.0.dist-info/METADATA

@@ -0,0 +1,92 @@
+Metadata-Version: 2.4
+Name: batchcorder-datafusion
+Version: 55.0.0
+Classifier: Programming Language :: Rust
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Topic :: Scientific/Engineering
+Requires-Dist: pyarrow>=18
+Summary: DataFusion table provider that tees one read-once Arrow stream to every scan in a query
+Author: Daniel Mesejo
+License: Apache-2.0
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Homepage, https://github.com/mesejo/batchcorder
+Project-URL: Repository, https://github.com/mesejo/batchcorder
+
+# batchcorder-datafusion
+
+A DataFusion table provider that tees one read-once Arrow stream to every
+scan in a query.
+
+Register a `TeeTableProvider` with a DataFusion `SessionContext`, and every
+scan of the table in a query (self-joins, CTEs referenced twice, subqueries)
+replays the same upstream stream, which is read only once. With the
+default `bounding="sealed"`, each planned scan reserves a reader at planning
+time, and each batch is evicted as soon as every scan has read past it.
+
+## Install
+
+The provider plugs into a DataFusion host through `datafusion-ffi`. It does
+not install a host: install the one you use, in a version from the
+compatibility table below.
+
+```bash
+pip install batchcorder-datafusion "datafusion>=55,<56"        # datafusion-python
+pip install batchcorder-datafusion "xorq-datafusion>=0.3,<0.4" # xorq-datafusion
+```
+
+## Usage
+
+```python
+import pyarrow as pa
+from datafusion import SessionContext
+
+from batchcorder_datafusion import TeeTableProvider
+
+
+schema = pa.schema([("k", pa.int64())])
+reader = pa.RecordBatchReader.from_batches(
+    schema, [pa.record_batch([[1, 2, 3]], schema=schema)]
+)
+
+ctx = SessionContext()
+ctx.register_table("t", TeeTableProvider(reader))
+ctx.sql("select * from t a join t b using (k)").collect()
+```
+
+The upstream can be any object that implements the Arrow PyCapsule stream
+interface (`__arrow_c_stream__`).
+
+A sealed provider serves one query. Register a new provider per query, or
+pass `bounding="unbounded"` for a table that is queried repeatedly.
+
+### Do not pass a `batchcorder.StreamCache`
+
+The provider keeps its own cache. Passing a `batchcorder` reader caches the
+data twice. Pass the original upstream reader, unless you also need to replay
+the stream outside DataFusion.
+
+## Compatibility
+
+`datafusion-ffi` structs are only compatible within one DataFusion major, so
+the provider and the host must be built on the same one. The major version of
+`batchcorder-datafusion` is that DataFusion major.
+
+| batchcorder-datafusion | DataFusion | datafusion (datafusion-python) | xorq-datafusion |
+|------------------------|------------|--------------------------------|-----------------|
+| 55.x                   | 55         | `>=55,<56`                     | `>=0.3,<0.4`    |
+
+Registering a `TeeTableProvider` with a host outside these ranges raises
+`RuntimeError`, naming the host version found and the range needed, before any
+data crosses the FFI boundary.
+
+`batchcorder-datafusion` does not depend on the `batchcorder` package. The
+cache is compiled into its extension, and data enters through the Arrow C
+stream interface.
+

batchcorder_datafusion-55.0.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+batchcorder_datafusion/__init__.py,sha256=iP9G6JsrPgzSa4pG-1qNRNSz360d0Ds40dQosinSiAg,1293
+batchcorder_datafusion/_batchcorder_datafusion.pyd,sha256=9V5k6wFh4bCdMhsWvCA8qLmVvv9Bq7dS0hkRIkZieGk,61756416
+batchcorder_datafusion/_batchcorder_datafusion.pyi,sha256=Rw2X2PUdtQVfQ9E3Z6HlL9yCs7Yc4bshPtvhlCCNpcg,975
+batchcorder_datafusion/_hosts.py,sha256=d7k1tnuA7E0v4faZbECWpsUHc-GQhHhQ2Zvwi4I7yCA,5508
+batchcorder_datafusion/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+batchcorder_datafusion-55.0.0.dist-info/METADATA,sha256=qqChFaj322NtiuVsZHPfvNA4IZzJ-kA5yF200fxWXMs,3590
+batchcorder_datafusion-55.0.0.dist-info/WHEEL,sha256=TzjyIyBUptJHnlQyOymQYLWjtXe7MGI0m3QvmcG7PpI,96
+batchcorder_datafusion-55.0.0.dist-info/sboms/batchcorder-datafusion-py.cyclonedx.json,sha256=_ZRMrakGPQW12EGJ3W2AtExVvjI28SD4Kda5ImubapY,345548
+batchcorder_datafusion-55.0.0.dist-info/RECORD,,

batchcorder_datafusion-55.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: maturin (1.16.0)
+Root-Is-Purelib: false
+Tag: cp310-abi3-win_amd64