cudf-polars-cu13 25.10.0__py3-none-any.whl

cudf_polars/experimental/base.py

@@ -0,0 +1,386 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Multi-partition base classes."""
+
+from __future__ import annotations
+
+import dataclasses
+from collections import defaultdict
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Generic, NamedTuple, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Iterator, MutableMapping
+
+    from cudf_polars.dsl.expr import NamedExpr
+    from cudf_polars.dsl.ir import IR
+    from cudf_polars.dsl.nodebase import Node
+
+
+class PartitionInfo:
+    """Partitioning information."""
+
+    __slots__ = ("count", "partitioned_on")
+    count: int
+    """Partition count."""
+    partitioned_on: tuple[NamedExpr, ...]
+    """Columns the data is hash-partitioned on."""
+
+    def __init__(
+        self,
+        count: int,
+        partitioned_on: tuple[NamedExpr, ...] = (),
+    ):
+        self.count = count
+        self.partitioned_on = partitioned_on
+
+    def keys(self, node: Node) -> Iterator[tuple[str, int]]:
+        """Return the partitioned keys for a given node."""
+        name = get_key_name(node)
+        yield from ((name, i) for i in range(self.count))
+
+    def __rich_repr__(self) -> Generator[Any, None, None]:
+        """Formatting for rich.pretty.pprint."""
+        yield "count", self.count
+        yield "partitioned_on", self.partitioned_on
+
+
+def get_key_name(node: Node) -> str:
+    """Generate the key name for a Node."""
+    return f"{type(node).__name__.lower()}-{hash(node)}"
+
+
+T = TypeVar("T")
+
+
+@dataclasses.dataclass
+class ColumnStat(Generic[T]):
+    """
+    Generic column-statistic.
+
+    Parameters
+    ----------
+    value
+        Statistics value. Value will be None
+        if the statistics is unknown.
+    exact
+        Whether the statistics is known exactly.
+    """
+
+    value: T | None = None
+    exact: bool = False
+
+
+@dataclasses.dataclass
+class UniqueStats:
+    """
+    Sampled unique-value statistics.
+
+    Parameters
+    ----------
+    count
+        Unique-value count.
+    fraction
+        Unique-value fraction. This corresponds to the total
+        number of unique values (count) divided by the total
+        number of rows.
+
+    Notes
+    -----
+    This class is used to track unique-value column statistics
+    that have been sampled from a data source.
+    """
+
+    count: ColumnStat[int] = dataclasses.field(default_factory=ColumnStat[int])
+    fraction: ColumnStat[float] = dataclasses.field(default_factory=ColumnStat[float])
+
+
+class DataSourceInfo:
+    """
+    Table data source information.
+
+    Notes
+    -----
+    This class should be sub-classed for specific
+    data source types (e.g. Parquet, DataFrame, etc.).
+    The required properties/methods enable lazy
+    sampling of the underlying datasource.
+    """
+
+    _unique_stats_columns: set[str]
+
+    @property
+    def row_count(self) -> ColumnStat[int]:  # pragma: no cover
+        """Data source row-count estimate."""
+        raise NotImplementedError("Sub-class must implement row_count.")
+
+    def unique_stats(self, column: str) -> UniqueStats:  # pragma: no cover
+        """Return unique-value statistics for a column."""
+        raise NotImplementedError("Sub-class must implement unique_stats.")
+
+    def storage_size(self, column: str) -> ColumnStat[int]:
+        """Return the average column size for a single file."""
+        return ColumnStat[int]()
+
+    @property
+    def unique_stats_columns(self) -> set[str]:
+        """Return the set of columns needing unique-value information."""
+        return self._unique_stats_columns
+
+    def add_unique_stats_column(self, column: str) -> None:
+        """Add a column needing unique-value information."""
+        self._unique_stats_columns.add(column)
+
+
+class DataSourcePair(NamedTuple):
+    """Pair of table-source and column-name information."""
+
+    table_source: DataSourceInfo
+    column_name: str
+
+
+class ColumnSourceInfo:
+    """
+    Source column information.
+
+    Parameters
+    ----------
+    table_source_pairs
+        Sequence of DataSourcePair objects.
+        Union operations will result in multiple elements.
+
+    Notes
+    -----
+    This is a thin wrapper around DataSourceInfo that provides
+    direct access to column-specific information.
+    """
+
+    __slots__ = (
+        "implied_unique_count",
+        "table_source_pairs",
+    )
+    table_source_pairs: list[DataSourcePair]
+    implied_unique_count: ColumnStat[int]
+    """Unique-value count implied by join heuristics."""
+
+    def __init__(self, *table_source_pairs: DataSourcePair) -> None:
+        self.table_source_pairs = list(table_source_pairs)
+        self.implied_unique_count = ColumnStat[int](None)
+
+    @property
+    def is_unique_stats_column(self) -> bool:
+        """Return whether this column requires unique-value information."""
+        return any(
+            pair.column_name in pair.table_source.unique_stats_columns
+            for pair in self.table_source_pairs
+        )
+
+    @property
+    def row_count(self) -> ColumnStat[int]:
+        """Data source row-count estimate."""
+        return ColumnStat[int](
+            # Use sum of table-source row-count estimates.
+            value=sum(
+                value
+                for pair in self.table_source_pairs
+                if (value := pair.table_source.row_count.value) is not None
+            )
+            or None,
+            # Row-count may be exact if there is only one table source.
+            exact=len(self.table_source_pairs) == 1
+            and self.table_source_pairs[0].table_source.row_count.exact,
+        )
+
+    def unique_stats(self, *, force: bool = False) -> UniqueStats:
+        """
+        Return unique-value statistics for a column.
+
+        Parameters
+        ----------
+        force
+            If True, return unique-value statistics even if the column
+            wasn't marked as needing unique-value information.
+        """
+        if (force or self.is_unique_stats_column) and len(self.table_source_pairs) == 1:
+            # Single table source.
+            # TODO: Handle multiple tables sources if/when necessary.
+            # We may never need to do this if the source unique-value
+            # statistics are only "used" by the Scan/DataFrameScan nodes.
+            table_source, column_name = self.table_source_pairs[0]
+            return table_source.unique_stats(column_name)
+        else:
+            # Avoid sampling unique-stats if this column
+            # wasn't marked as "needing" unique-stats.
+            return UniqueStats()
+
+    @property
+    def storage_size(self) -> ColumnStat[int]:
+        """Return the average column size for a single file."""
+        # We don't need to handle concatenated statistics for ``storage_size``.
+        # Just return the storage size of the first table source.
+        if self.table_source_pairs:
+            table_source, column_name = self.table_source_pairs[0]
+            return table_source.storage_size(column_name)
+        else:  # pragma: no cover; We never call this for empty table sources.
+            return ColumnStat[int]()
+
+    def add_unique_stats_column(self, column: str | None = None) -> None:
+        """Add a column needing unique-value information."""
+        # We must call add_unique_stats_column for ALL table sources.
+        for table_source, column_name in self.table_source_pairs:
+            table_source.add_unique_stats_column(column or column_name)
+
+
+class ColumnStats:
+    """
+    Column statistics.
+
+    Parameters
+    ----------
+    name
+        Column name.
+    children
+        Child ColumnStats objects.
+    source_info
+        Column source information.
+    unique_count
+        Unique-value count.
+    """
+
+    __slots__ = ("children", "name", "source_info", "unique_count")
+
+    name: str
+    children: tuple[ColumnStats, ...]
+    source_info: ColumnSourceInfo
+    unique_count: ColumnStat[int]
+
+    def __init__(
+        self,
+        name: str,
+        *,
+        children: tuple[ColumnStats, ...] = (),
+        source_info: ColumnSourceInfo | None = None,
+        unique_count: ColumnStat[int] | None = None,
+    ) -> None:
+        self.name = name
+        self.children = children
+        self.source_info = source_info or ColumnSourceInfo()
+        self.unique_count = unique_count or ColumnStat[int](None)
+
+    def new_parent(
+        self,
+        *,
+        name: str | None = None,
+    ) -> ColumnStats:
+        """
+        Initialize a new parent ColumnStats object.
+
+        Parameters
+        ----------
+        name
+            The new column name.
+
+        Returns
+        -------
+        A new ColumnStats object.
+
+        Notes
+        -----
+        This API preserves the original DataSourceInfo reference.
+        """
+        return ColumnStats(
+            name=name or self.name,
+            children=(self,),
+            # Want to reference the same DataSourceInfo
+            source_info=self.source_info,
+        )
+
+
+class JoinKey:
+    """
+    Join-key information.
+
+    Parameters
+    ----------
+    column_stats
+        Column statistics for the join key.
+
+    Notes
+    -----
+    This class is used to track join-key information.
+    It is used to track the columns being joined on
+    and the estimated unique-value count for the join key.
+    """
+
+    column_stats: tuple[ColumnStats, ...]
+    implied_unique_count: int | None
+    """Estimated unique-value count from join heuristics."""
+
+    def __init__(self, *column_stats: ColumnStats) -> None:
+        self.column_stats = column_stats
+        self.implied_unique_count = None
+
+    @cached_property
+    def source_row_count(self) -> int | None:
+        """
+        Return the estimated row-count of the source columns.
+
+        Notes
+        -----
+        This is the maximum row-count estimate of the source columns.
+        """
+        return max(
+            (
+                cs.source_info.row_count.value
+                for cs in self.column_stats
+                if cs.source_info.row_count.value is not None
+            ),
+            default=None,
+        )
+
+
+class JoinInfo:
+    """
+    Join information.
+
+    Notes
+    -----
+    This class is used to track mappings between joined-on
+    columns and joined-on keys (groups of columns). We need
+    these mappings to calculate equivalence sets and make
+    join-based unique-count and row-count estimates.
+    """
+
+    __slots__ = ("column_map", "join_map", "key_map")
+
+    column_map: MutableMapping[ColumnStats, set[ColumnStats]]
+    """Mapping between joined columns."""
+    key_map: MutableMapping[JoinKey, set[JoinKey]]
+    """Mapping between joined keys (groups of columns)."""
+    join_map: dict[IR, list[JoinKey]]
+    """Mapping between IR nodes and associated join keys."""
+
+    def __init__(self) -> None:
+        self.column_map: MutableMapping[ColumnStats, set[ColumnStats]] = defaultdict(
+            set[ColumnStats]
+        )
+        self.key_map: MutableMapping[JoinKey, set[JoinKey]] = defaultdict(set[JoinKey])
+        self.join_map: dict[IR, list[JoinKey]] = {}
+
+
+class StatsCollector:
+    """Column statistics collector."""
+
+    __slots__ = ("column_stats", "join_info", "row_count")
+
+    row_count: dict[IR, ColumnStat[int]]
+    """Estimated row count for each IR node."""
+    column_stats: dict[IR, dict[str, ColumnStats]]
+    """Column statistics for each IR node."""
+    join_info: JoinInfo
+    """Join information."""
+
+    def __init__(self) -> None:
+        self.row_count: dict[IR, ColumnStat[int]] = {}
+        self.column_stats: dict[IR, dict[str, ColumnStats]] = {}
+        self.join_info = JoinInfo()

cudf_polars/experimental/benchmarks/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Experimental benchmarks."""

cudf_polars/experimental/benchmarks/pdsds.py

@@ -0,0 +1,220 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Experimental PDS-DS benchmarks.
+
+Based on https://github.com/pola-rs/polars-benchmark.
+
+WARNING: This is an experimental (and unofficial)
+benchmark script. It is not intended for public use
+and may be modified or removed at any time.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import importlib
+import os
+import time
+from collections import defaultdict
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import polars as pl
+
+with contextlib.suppress(ImportError):
+    from cudf_polars.experimental.benchmarks.utils import (
+        Record,
+        RunConfig,
+        get_executor_options,
+        parse_args,
+        run_polars,
+    )
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from types import ModuleType
+    from typing import Any
+
+# Without this setting, the first IO task to run
+# on each worker takes ~15 sec extra
+os.environ["KVIKIO_COMPAT_MODE"] = os.environ.get("KVIKIO_COMPAT_MODE", "on")
+os.environ["KVIKIO_NTHREADS"] = os.environ.get("KVIKIO_NTHREADS", "8")
+
+
+def valid_query(name: str) -> bool:
+    """Return True for valid query names eg. 'q9', 'q65', etc."""
+    if not name.startswith("q"):
+        return False
+    try:
+        q_num = int(name[1:])
+    except ValueError:
+        return False
+    else:
+        return 1 <= q_num <= 99
+
+
+class PDSDSQueriesMeta(type):
+    """Metaclass used for query lookup."""
+
+    def __getattr__(cls, name: str):  # type: ignore
+        """Query lookup."""
+        if valid_query(name):
+            q_num = int(name[1:])
+            module: ModuleType = importlib.import_module(
+                f"cudf_polars.experimental.benchmarks.pdsds_queries.q{q_num}"
+            )
+            return getattr(module, cls.q_impl)
+        raise AttributeError(f"{name} is not a valid query name")
+
+
+class PDSDSQueries(metaclass=PDSDSQueriesMeta):
+    """Base class for query loading."""
+
+    q_impl: str
+    name: str = "pdsds"
+
+
+class PDSDSPolarsQueries(PDSDSQueries):
+    """Polars Queries."""
+
+    q_impl = "polars_impl"
+
+
+class PDSDSDuckDBQueries(PDSDSQueries):
+    """DuckDB Queries."""
+
+    q_impl = "duckdb_impl"
+
+
+def execute_duckdb_query(query: str, dataset_path: Path) -> pl.DataFrame:
+    """Execute a query with DuckDB."""
+    import duckdb
+
+    conn = duckdb.connect()
+
+    statements = [
+        f"CREATE VIEW {table.stem} as SELECT * FROM read_parquet('{table.absolute()}');"
+        for table in Path(dataset_path).glob("*.parquet")
+    ]
+    statements.append(query)
+    return conn.execute("\n".join(statements)).pl()
+
+
+def run_duckdb(benchmark: Any, options: Sequence[str] | None = None) -> None:
+    """Run the benchmark with DuckDB."""
+    args = parse_args(options, num_queries=99)
+    vars(args).update({"query_set": benchmark.name})
+    run_config = RunConfig.from_args(args)
+    records: defaultdict[int, list[Record]] = defaultdict(list)
+
+    for q_id in run_config.queries:
+        try:
+            duckdb_query = getattr(PDSDSDuckDBQueries, f"q{q_id}")(run_config)
+        except AttributeError as err:
+            raise NotImplementedError(f"Query {q_id} not implemented.") from err
+
+        print(f"DuckDB Executing: {q_id}")
+        records[q_id] = []
+
+        for i in range(args.iterations):
+            t0 = time.time()
+
+            result = execute_duckdb_query(duckdb_query, run_config.dataset_path)
+
+            t1 = time.time()
+            record = Record(query=q_id, duration=t1 - t0)
+            if args.print_results:
+                print(result)
+
+            print(f"Query {q_id} - Iteration {i} finished in {record.duration:0.4f}s")
+            records[q_id].append(record)
+
+
+def run_validate(benchmark: Any, options: Sequence[str] | None = None) -> None:
+    """Validate Polars CPU vs DuckDB or Polars GPU."""
+    from polars.testing import assert_frame_equal
+
+    args = parse_args(options, num_queries=99)
+    vars(args).update({"query_set": benchmark.name})
+    run_config = RunConfig.from_args(args)
+
+    baseline = args.baseline
+    if baseline not in {"duckdb", "cpu"}:
+        raise ValueError("Baseline must be one of: 'duckdb', 'cpu'")
+
+    failures: list[int] = []
+
+    engine: pl.GPUEngine | None = None
+    if run_config.executor != "cpu":
+        engine = pl.GPUEngine(
+            raise_on_fail=True,
+            executor=run_config.executor,
+            executor_options=get_executor_options(run_config, PDSDSPolarsQueries),
+        )
+
+    for q_id in run_config.queries:
+        print(f"\nValidating Query {q_id}")
+        try:
+            polars_query = getattr(PDSDSPolarsQueries, f"q{q_id}")(run_config)
+            duckdb_query = getattr(PDSDSDuckDBQueries, f"q{q_id}")(run_config)
+        except AttributeError as err:
+            raise NotImplementedError(f"Query {q_id} not implemented.") from err
+
+        if baseline == "duckdb":
+            base_result = execute_duckdb_query(duckdb_query, run_config.dataset_path)
+        elif baseline == "cpu":
+            base_result = polars_query.collect(new_streaming=True)
+
+        if run_config.executor == "cpu":
+            test_result = polars_query.collect(new_streaming=True)
+        else:
+            try:
+                test_result = polars_query.collect(engine=engine)
+            except Exception as e:
+                failures.append(q_id)
+                print(f"❌ Query {q_id} failed validation: GPU execution failed.\n{e}")
+                continue
+
+        try:
+            assert_frame_equal(
+                base_result,
+                test_result,
+                check_dtypes=True,
+                check_column_order=False,
+            )
+            print(f"✅ Query {q_id} passed validation.")
+        except AssertionError as e:
+            failures.append(q_id)
+            print(f"❌ Query {q_id} failed validation:\n{e}")
+            if args.print_results:
+                print("Baseline Result:\n", base_result)
+                print("Test Result:\n", test_result)
+
+    if failures:
+        print("\nValidation Summary:")
+        print("===================")
+        print(f"{len(failures)} query(s) failed: {failures}")
+    else:
+        print("\nAll queries passed validation.")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Run PDS-DS benchmarks.")
+    parser.add_argument(
+        "--engine",
+        choices=["polars", "duckdb", "validate"],
+        default="polars",
+        help="Which engine to use for executing the benchmarks or to validate results.",
+    )
+    args, extra_args = parser.parse_known_args()
+
+    if args.engine == "polars":
+        run_polars(PDSDSPolarsQueries, extra_args, num_queries=99)
+    elif args.engine == "duckdb":
+        run_duckdb(PDSDSDuckDBQueries, extra_args)
+    elif args.engine == "validate":
+        run_validate(PDSDSQueries, extra_args)

cudf_polars/experimental/benchmarks/pdsds_queries/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""DuckDB and Polars queries."""

cudf_polars/experimental/benchmarks/pdsds_queries/q1.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Query 1."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import polars as pl
+
+from cudf_polars.experimental.benchmarks.utils import get_data
+
+if TYPE_CHECKING:
+    from cudf_polars.experimental.benchmarks.utils import RunConfig
+
+
+def duckdb_impl(run_config: RunConfig) -> str:
+    """Query 1."""
+    return """
+    WITH customer_total_return
+        AS (SELECT sr_customer_sk     AS ctr_customer_sk,
+                    sr_store_sk        AS ctr_store_sk,
+                    Sum(sr_return_amt) AS ctr_total_return
+            FROM   store_returns,
+                    date_dim
+            WHERE  sr_returned_date_sk = d_date_sk
+                    AND d_year = 2001
+            GROUP  BY sr_customer_sk,
+                    sr_store_sk)
+    SELECT c_customer_id
+    FROM   customer_total_return ctr1,
+        store,
+        customer
+    WHERE  ctr1.ctr_total_return > (SELECT Avg(ctr_total_return) * 1.2
+                                    FROM   customer_total_return ctr2
+                                    WHERE  ctr1.ctr_store_sk = ctr2.ctr_store_sk)
+        AND s_store_sk = ctr1.ctr_store_sk
+        AND s_state = 'TN'
+        AND ctr1.ctr_customer_sk = c_customer_sk
+    ORDER  BY c_customer_id
+    LIMIT 100;
+    """
+
+
+def polars_impl(run_config: RunConfig) -> pl.LazyFrame:
+    """Query 1."""
+    store_returns = get_data(
+        run_config.dataset_path, "store_returns", run_config.suffix
+    )
+    date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix)
+    store = get_data(run_config.dataset_path, "store", run_config.suffix)
+    customer = get_data(run_config.dataset_path, "customer", run_config.suffix)
+
+    # Step 1: Create customer_total_return CTE equivalent
+    customer_total_return = (
+        store_returns.join(
+            date_dim, left_on="sr_returned_date_sk", right_on="d_date_sk"
+        )
+        .filter(pl.col("d_year") == 2001)
+        .group_by(["sr_customer_sk", "sr_store_sk"])
+        .agg(pl.col("sr_return_amt").sum().alias("ctr_total_return"))
+        .rename(
+            {
+                "sr_customer_sk": "ctr_customer_sk",
+                "sr_store_sk": "ctr_store_sk",
+            }
+        )
+    )
+
+    # Step 2: Calculate average return per store for the subquery
+    store_avg_returns = customer_total_return.group_by("ctr_store_sk").agg(
+        [(pl.col("ctr_total_return").mean() * 1.2).alias("avg_return_threshold")]
+    )
+
+    # Step 3: Join everything together and apply filters
+    return (
+        customer_total_return.join(
+            store_avg_returns, left_on="ctr_store_sk", right_on="ctr_store_sk"
+        )
+        .filter(pl.col("ctr_total_return") > pl.col("avg_return_threshold"))
+        .join(store, left_on="ctr_store_sk", right_on="s_store_sk")
+        .filter(pl.col("s_state") == "TN")
+        .join(customer, left_on="ctr_customer_sk", right_on="c_customer_sk")
+        .select(["c_customer_id"])
+        .sort("c_customer_id")
+        .limit(100)
+    )