vast-daft 0.1.1__tar.gz

vast_daft-0.1.1/.gitignore

@@ -0,0 +1,43 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.egg-info/
+*.egg
+dist/
+build/
+.eggs/
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# uv
+uv.lock
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+.DS_Store
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.mypy_cache/
+
+# Distribution
+*.whl
+*.tar.gz
+
+# Environment
+.env
+.env.*
+!.env.example
+
+# Local example data (Lance datasets, downloaded corpora)
+data/

vast_daft-0.1.1/PKG-INFO

@@ -0,0 +1,82 @@
+Metadata-Version: 2.4
+Name: vast-daft
+Version: 0.1.1
+Summary: Daft custom connector (DataSource/DataSink) for VastDB
+Project-URL: Homepage, https://github.com/vast-data/vast-daft
+Project-URL: Repository, https://github.com/vast-data/vast-daft
+Author-email: Ofer Helman <ofer.helman@vastdata.com>
+License-Expression: Apache-2.0
+Requires-Python: >=3.12
+Requires-Dist: daft==0.7.10
+Requires-Dist: ibis-framework>=9.0
+Requires-Dist: numpy<2
+Requires-Dist: pyarrow>=15.0
+Requires-Dist: vastdb>=1.2
+Provides-Extra: dev
+Requires-Dist: pyiceberg[s3fs,sql-sqlite]>=0.11.1; extra == 'dev'
+Requires-Dist: pylance>=0.39.0; extra == 'dev'
+Requires-Dist: pyright>=1.1; extra == 'dev'
+Requires-Dist: pytest-cov>=5.0; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: python-dotenv>=1.0; extra == 'dev'
+Requires-Dist: ray==2.53.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Provides-Extra: examples
+Requires-Dist: python-dotenv>=1.0; extra == 'examples'
+Description-Content-Type: text/markdown
+
+# vast-daft
+
+Daft custom connector (`DataSource` / `DataSink`) for [VastDB](https://vastdata.com).
+
+## Installation
+
+```bash
+pip install vast-daft
+```
+
+Or with uv:
+
+```bash
+uv add vast-daft
+```
+
+## Quick Start
+
+### Reading from VastDB
+
+```python
+import pyarrow as pa
+from vast_daft import VastDBConfig, VastDBDataSource
+
+config = VastDBConfig(
+    endpoint="http://vastdb:9090",
+    access_key="YOUR_ACCESS_KEY",
+    secret_key="YOUR_SECRET_KEY",
+    bucket="my-bucket",
+    schema="my-schema",
+)
+
+schema = pa.schema([
+    ("id", pa.string()),
+    ("name", pa.string()),
+    ("value", pa.float64()),
+])
+
+source = VastDBDataSource(config, "my_table", schema)
+df = source.read()
+df.show()
+```
+
+### Writing to VastDB
+
+```python
+import daft
+from vast_daft import VastDBConfig, VastDBDataSink
+
+config = VastDBConfig(...)
+schema = pa.schema([("id", pa.string()), ("value", pa.float64())])
+
+sink = VastDBDataSink(config, "my_table", schema)
+daft.from_pydict({"id": ["a", "b"], "value": [1.0, 2.0]}).write_sink(sink).show()
+```

vast_daft-0.1.1/README.md

@@ -0,0 +1,55 @@
+# vast-daft
+
+Daft custom connector (`DataSource` / `DataSink`) for [VastDB](https://vastdata.com).
+
+## Installation
+
+```bash
+pip install vast-daft
+```
+
+Or with uv:
+
+```bash
+uv add vast-daft
+```
+
+## Quick Start
+
+### Reading from VastDB
+
+```python
+import pyarrow as pa
+from vast_daft import VastDBConfig, VastDBDataSource
+
+config = VastDBConfig(
+    endpoint="http://vastdb:9090",
+    access_key="YOUR_ACCESS_KEY",
+    secret_key="YOUR_SECRET_KEY",
+    bucket="my-bucket",
+    schema="my-schema",
+)
+
+schema = pa.schema([
+    ("id", pa.string()),
+    ("name", pa.string()),
+    ("value", pa.float64()),
+])
+
+source = VastDBDataSource(config, "my_table", schema)
+df = source.read()
+df.show()
+```
+
+### Writing to VastDB
+
+```python
+import daft
+from vast_daft import VastDBConfig, VastDBDataSink
+
+config = VastDBConfig(...)
+schema = pa.schema([("id", pa.string()), ("value", pa.float64())])
+
+sink = VastDBDataSink(config, "my_table", schema)
+daft.from_pydict({"id": ["a", "b"], "value": [1.0, 2.0]}).write_sink(sink).show()
+```

vast_daft-0.1.1/pyproject.toml

@@ -0,0 +1,59 @@
+[project]
+name = "vast-daft"
+version = "0.1.1"
+description = "Daft custom connector (DataSource/DataSink) for VastDB"
+readme = "README.md"
+requires-python = ">=3.12"
+license = "Apache-2.0"
+authors = [{ name = "Ofer Helman", email = "ofer.helman@vastdata.com" }]
+dependencies = [
+  "daft==0.7.10",
+  "numpy<2",
+  "vastdb>=1.2",
+  "pyarrow>=15.0",
+  "ibis-framework>=9.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/vast-data/vast-daft"
+Repository = "https://github.com/vast-data/vast-daft"
+
+[project.optional-dependencies]
+dev = [
+  "pytest>=8.0",
+  "pytest-cov>=5.0",
+  "ruff>=0.4",
+  "pyright>=1.1",
+  # "ty>=0.0.0a1",
+  "python-dotenv>=1.0",
+  "ray==2.53.0",
+  "pylance>=0.39.0",
+  "pyiceberg[s3fs,sql-sqlite]>=0.11.1",
+]
+examples = [
+  "python-dotenv>=1.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.pyright]
+typeCheckingMode = "basic"
+
+[tool.pytest]
+testpaths = ["tests"]
+
+[tool.ruff]
+target-version = "py312"
+line-length = 120
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "W", "UP"]
+
+[tool.ty.environment]
+python-version = "3.12"
+
+[tool.ty.rules]
+unsupported-operator = "ignore"
+

vast_daft-0.1.1/src/vast_daft/__init__.py

@@ -0,0 +1,186 @@
+"""vast-daft: Daft custom connector for VastDB.
+
+Provides ``DataSource`` / ``DataSink`` implementations for reading from
+and writing to VastDB, and a Daft catalog integration for native SQL queries.
+
+Importing this module also registers two convenience methods on
+``daft.DataFrame``:
+
+* ``df.read_vastdb(config, table_name, ...)`` — read a VastDB table
+* ``df.write_vastdb(config, table_name, schema, ...)`` — write to a VastDB table
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import daft
+import pyarrow as pa
+
+from vast_daft.config import VastDBConfig
+from vast_daft.connection import VastDBConnection
+from vast_daft.scan import VastDBScanOperator
+from vast_daft.sink import VastDBDataSink
+from vast_daft.source import VastDBDataSource
+from vast_daft.table import VastDBCatalog, VastDBTable
+
+# ---------------------------------------------------------------------------
+# df.read_vastdb()
+# ---------------------------------------------------------------------------
+
+def _read_vastdb(
+    config: VastDBConfig,
+    table_name: str,
+    *,
+    bucket: str | None = None,
+    schema: str | None = None,
+    columns: list[str] | None = None,
+    num_splits: int | None = None,
+) -> daft.DataFrame:
+    """Read a VastDB table into a Daft DataFrame.
+
+    This is a module-level factory (not a DataFrame method) that mirrors the
+    ergonomics of ``daft.read_parquet()``, ``daft.read_iceberg()``, etc.
+
+    Parameters
+    ----------
+    config:
+        VastDB connection configuration.
+    table_name:
+        Name of the VastDB table to read.
+    bucket:
+        VastDB bucket.  Overrides ``config.bucket`` when provided.
+    schema:
+        VastDB schema.  Overrides ``config.schema`` when provided.
+    columns:
+        Optional list of column names to project.
+    num_splits:
+        Number of parallel read splits.  Auto-estimated when omitted.
+
+    Returns
+    -------
+    daft.DataFrame
+
+    Examples
+    --------
+    >>> import vast_daft
+    >>> from vast_daft import VastDBConfig
+    >>> config = VastDBConfig(endpoint="http://vastdb:9090", access_key="...", secret_key="...")
+    >>> df = vast_daft.read_vastdb(config, "orders")
+    >>> df.show()
+    """
+    from daft.daft import ScanOperatorHandle
+    from daft.logical.builder import LogicalPlanBuilder
+
+    _bucket = bucket or config.bucket
+    _schema = schema or config.schema
+    if _bucket is None or _schema is None:
+        raise ValueError("bucket and schema must be provided either as arguments or via VastDBConfig")
+
+    # Discover the PyArrow schema from the live table
+    conn = VastDBConnection(config)
+    with conn.session.transaction() as tx:
+        db_table = tx.bucket(_bucket).schema(_schema).table(table_name)
+        pa_schema = db_table.columns()
+
+    if columns is not None:
+        pa_schema = pa.schema([pa_schema.field(c) for c in columns])
+
+    kwargs: dict[str, Any] = {}
+    if num_splits is not None:
+        kwargs["num_splits"] = num_splits
+
+    scan_op = VastDBScanOperator(
+        config,
+        table_name,
+        pa_schema,
+        bucket=_bucket,
+        db_schema=_schema,
+        columns=columns,
+        **kwargs,
+    )
+    handle = ScanOperatorHandle.from_python_scan_operator(scan_op)
+    builder = LogicalPlanBuilder.from_tabular_scan(scan_operator=handle)
+    return daft.DataFrame(builder)
+
+
+# ---------------------------------------------------------------------------
+# df.write_vastdb()
+# ---------------------------------------------------------------------------
+
+def _write_vastdb(
+    self: daft.DataFrame,
+    config: VastDBConfig,
+    table_name: str,
+    table_schema: pa.Schema,
+    *,
+    bucket: str | None = None,
+    schema: str | None = None,
+    create_if_missing: bool = True,
+) -> daft.DataFrame:
+    """Write the DataFrame to a VastDB table.
+
+    Registered on ``daft.DataFrame`` when ``vast_daft`` is imported.
+
+    Parameters
+    ----------
+    config:
+        VastDB connection configuration.
+    table_name:
+        Name of the target VastDB table.
+    table_schema:
+        PyArrow schema of the target table.
+    bucket:
+        VastDB bucket.  Overrides ``config.bucket`` when provided.
+    schema:
+        VastDB schema.  Overrides ``config.schema`` when provided.
+    create_if_missing:
+        Auto-create the table if it does not exist (default: ``True``).
+
+    Returns
+    -------
+    daft.DataFrame
+        A DataFrame with ``rows_written`` and ``bytes_written`` columns,
+        one row per partition written.
+
+    Examples
+    --------
+    >>> import daft, pyarrow as pa, vast_daft
+    >>> from vast_daft import VastDBConfig
+    >>> config = VastDBConfig(endpoint="http://vastdb:9090", access_key="...", secret_key="...")
+    >>> schema = pa.schema([("id", pa.int64()), ("value", pa.float64())])
+    >>> df = daft.from_pydict({"id": [1, 2], "value": [0.1, 0.2]})
+    >>> df.write_vastdb(config, "my_table", schema)
+    """
+    sink = VastDBDataSink(
+        config,
+        table_name,
+        table_schema,
+        bucket=bucket,
+        schema=schema,
+        create_if_missing=create_if_missing,
+    )
+    return self.write_sink(sink)
+
+
+# Register df.write_vastdb on daft.DataFrame
+daft.DataFrame.write_vastdb = _write_vastdb  # type: ignore[attr-defined]
+
+
+__all__ = [
+    # Core
+    "VastDBConfig",
+    "VastDBConnection",
+    # Daft connectors
+    "VastDBDataSource",
+    "VastDBDataSink",
+    "VastDBScanOperator",
+    # Catalog / Table (Daft interfaces — use with daft.attach_catalog + daft.sql)
+    "VastDBCatalog",
+    "VastDBTable",
+    # Convenience top-level functions (mirrors daft.read_parquet etc.)
+    "read_vastdb",
+]
+
+# Public alias — lets users call vast_daft.read_vastdb(...)
+read_vastdb = _read_vastdb

vast_daft-0.1.1/src/vast_daft/_pushdown.py

@@ -0,0 +1,245 @@
+"""Translate Daft ``Pushdowns`` to ibis predicates and projection columns.
+
+Daft calls ``get_tasks(pushdowns)`` with a :class:`~daft.io.pushdowns.Pushdowns`
+object that may carry:
+
+* ``filters``  — a ``daft.expressions.Expression`` (the combined filter predicate)
+* ``columns``  — a ``list[str]`` of column names to project
+* ``limit``    — an ``int`` row limit
+
+This module converts the ``filters`` expression into an ibis predicate that can
+be passed to :func:`vastdb.table.Table.select_splits`, achieving genuine
+server-side push-down for filter, column projection, and row limit.
+
+Supported Daft expression nodes
+---------------------------------
+Binary comparisons: ``==``, ``!=``, ``>``, ``>=``, ``<``, ``<=``
+Membership:         ``col.is_in([...])``
+Between:            ``col.between(lo, hi)``
+Null checks:        ``col.is_null()``, ``col.not_null()``
+Logical:            ``expr & expr``, ``expr | expr``, ``~expr``
+String:             ``col.contains(s)``, ``col.startswith(s)``, ``col.endswith(s)``
+
+Any unsupported node causes the entire filter to be returned un-translated
+(``None``), leaving Daft to apply it client-side as a post-filter.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import ibis
+from daft.expressions import Expression, ExpressionVisitor
+from daft.io.pushdowns import Pushdowns
+from ibis.common.deferred import Deferred
+from ibis.expr.types import Column as BooleanColumn
+
+logger = logging.getLogger(__name__)
+
+Predicate = BooleanColumn | Deferred
+
+# Sentinel to signal "unsupported — fall back to client-side"
+_UNSUPPORTED = object()
+
+
+# ---------------------------------------------------------------------------
+# Visitor
+# ---------------------------------------------------------------------------
+
+
+class _DaftToIbisVisitor(ExpressionVisitor):
+    """Walk a Daft expression tree and produce an ibis predicate.
+
+    Returns ``_UNSUPPORTED`` (the module-level sentinel) for any node that
+    cannot be translated, so callers can detect a partial failure without
+    raising.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.referenced_columns: set[str] = set()
+
+    # -- Leaves --------------------------------------------------------------
+
+    def visit_col(self, name: str) -> Any:
+        self.referenced_columns.add(name)
+        return ibis._[name]
+
+    def visit_lit(self, value: Any) -> Any:
+        # Daft passes the Python value directly (int, float, str, bool, None)
+        return value
+
+    # -- Structural nodes that we don't push down ----------------------------
+
+    def visit_alias(self, expr: Expression, alias: str) -> Any:
+        # Aliases don't affect predicate semantics — visit the inner expr
+        return self.visit(expr)
+
+    def visit_cast(self, expr: Expression, dtype: Any) -> Any:
+        # Casts are not translatable to ibis predicates
+        return _UNSUPPORTED
+
+    def visit_coalesce(self, *args: Expression) -> Any:
+        # Forward-compatibility with newer Daft visitor APIs.
+        return _UNSUPPORTED
+
+    # -- Function dispatch ---------------------------------------------------
+
+    def visit_function(self, name: str, args: list) -> Any:
+        """Fallback for unrecognised function names."""
+        return _UNSUPPORTED
+
+    # visit_() in the base class dispatches to visit_<name> if it exists,
+    # otherwise calls visit_function().  We register individual functions
+    # below as visit_<daft_function_name> methods.
+
+    def _binop(self, args: list, op: str) -> Any:
+        """Helper: visit two expression args and apply a binary ibis operator."""
+        lhs = self.visit(args[0])
+        rhs = self.visit(args[1])
+        if lhs is _UNSUPPORTED or rhs is _UNSUPPORTED:
+            return _UNSUPPORTED
+        ops = {
+            "equal": lambda a, b: a == b,
+            "not_equal": lambda a, b: a != b,
+            "greater_than": lambda a, b: a > b,
+            "greater_than_or_equal": lambda a, b: a >= b,
+            "less_than": lambda a, b: a < b,
+            "less_than_or_equal": lambda a, b: a <= b,
+            "and": lambda a, b: a & b,
+            "or": lambda a, b: a | b,
+        }
+        return ops[op](lhs, rhs)
+
+    # Comparison operators
+    def visit_equal(self, lhs: Expression, rhs: Expression) -> Any:
+        return self._binop([lhs, rhs], "equal")
+
+    def visit_not_equal(self, lhs: Expression, rhs: Expression) -> Any:
+        return self._binop([lhs, rhs], "not_equal")
+
+    def visit_greater_than(self, lhs: Expression, rhs: Expression) -> Any:
+        return self._binop([lhs, rhs], "greater_than")
+
+    def visit_greater_than_or_equal(self, lhs: Expression, rhs: Expression) -> Any:
+        return self._binop([lhs, rhs], "greater_than_or_equal")
+
+    def visit_less_than(self, lhs: Expression, rhs: Expression) -> Any:
+        return self._binop([lhs, rhs], "less_than")
+
+    def visit_less_than_or_equal(self, lhs: Expression, rhs: Expression) -> Any:
+        return self._binop([lhs, rhs], "less_than_or_equal")
+
+    # Logical operators
+    def visit_and(self, lhs: Expression, rhs: Expression) -> Any:
+        return self._binop([lhs, rhs], "and")
+
+    def visit_or(self, lhs: Expression, rhs: Expression) -> Any:
+        return self._binop([lhs, rhs], "or")
+
+    def visit_not(self, expr: Expression) -> Any:
+        inner = self.visit(expr)
+        if inner is _UNSUPPORTED:
+            return _UNSUPPORTED
+        return ~inner
+
+    # Null checks
+    def visit_is_null(self, expr: Expression) -> Any:
+        inner = self.visit(expr)
+        if inner is _UNSUPPORTED:
+            return _UNSUPPORTED
+        return inner.isnull()
+
+    def visit_not_null(self, expr: Expression) -> Any:
+        inner = self.visit(expr)
+        if inner is _UNSUPPORTED:
+            return _UNSUPPORTED
+        return inner.notnull()
+
+    # Membership: is_in(expr, [lit, lit, ...])
+    # NOTE: args[1] is a raw list of Expression objects, not pre-visited.
+    def visit_is_in(self, expr: Expression, values: list) -> Any:
+        col_ibis = self.visit(expr)
+        if col_ibis is _UNSUPPORTED:
+            return _UNSUPPORTED
+        visited_vals = [self.visit(v) for v in values]
+        if any(v is _UNSUPPORTED for v in visited_vals):
+            return _UNSUPPORTED
+        return col_ibis.isin(visited_vals)
+
+    # Between: between(expr, lo, hi)
+    def visit_between(self, expr: Expression, lo: Expression, hi: Expression) -> Any:
+        col_ibis = self.visit(expr)
+        lo_val = self.visit(lo)
+        hi_val = self.visit(hi)
+        if any(v is _UNSUPPORTED for v in [col_ibis, lo_val, hi_val]):
+            return _UNSUPPORTED
+        return col_ibis.between(lo_val, hi_val)
+
+    # String functions
+    def visit_starts_with(self, expr: Expression, prefix: Expression) -> Any:
+        col_ibis = self.visit(expr)
+        val = self.visit(prefix)
+        if col_ibis is _UNSUPPORTED or val is _UNSUPPORTED:
+            return _UNSUPPORTED
+        return col_ibis.startswith(val)
+
+    def visit_ends_with(self, expr: Expression, suffix: Expression) -> Any:
+        col_ibis = self.visit(expr)
+        val = self.visit(suffix)
+        if col_ibis is _UNSUPPORTED or val is _UNSUPPORTED:
+            return _UNSUPPORTED
+        return col_ibis.endswith(val)
+
+    def visit_contains(self, expr: Expression, pattern: Expression) -> Any:
+        col_ibis = self.visit(expr)
+        val = self.visit(pattern)
+        if col_ibis is _UNSUPPORTED or val is _UNSUPPORTED:
+            return _UNSUPPORTED
+        return col_ibis.contains(val)
+
+
+# ---------------------------------------------------------------------------
+# Public entry-point
+# ---------------------------------------------------------------------------
+
+def pushdowns_to_predicate(pushdowns: Pushdowns) -> Predicate | None:
+    """Convert *pushdowns.filters* to an ibis predicate, or ``None``.
+
+    Returns ``None`` when no filter is present or when the expression contains
+    an unsupported node.  In that case the caller should not pass a predicate
+    to VastDB, and Daft will apply the filter client-side.
+    """
+    predicate, _ = pushdowns_to_predicate_and_columns(pushdowns)
+    return predicate
+
+
+def pushdowns_to_predicate_and_columns(
+    pushdowns: Pushdowns,
+) -> tuple[Predicate | None, set[str]]:
+    """Translate the filter and return the columns it references.
+
+    The referenced column set is needed by callers that also push down a
+    column projection. VastDB's ``select_splits`` requires every column the
+    predicate touches to appear in ``columns``; otherwise the server raises
+    ``FieldNotFound``. Callers must union this set with their projection
+    before handing it to VastDB.
+
+    The set is empty when the predicate is unsupported or absent.
+    """
+    if pushdowns.filters is None:
+        return None, set()
+    visitor = _DaftToIbisVisitor()
+    try:
+        result = visitor.visit(pushdowns.filters)
+    except Exception as exc:
+        logger.debug("Daft filter not pushed down to VastDB — visitor error: %s", exc)
+        return None, set()
+    if result is _UNSUPPORTED:
+        logger.debug(
+            "Daft filter not pushed down to VastDB — unsupported expression: %s",
+            pushdowns.filters,
+        )
+        return None, set()
+    return result, visitor.referenced_columns