benchbox 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

benchbox/platforms/dataframe/polars_maintenance.py

@@ -0,0 +1,630 @@
+"""Polars Maintenance Operations Implementation.
+
+This module implements DataFrame maintenance operations for Polars,
+enabling TPC-H RF1/RF2 and TPC-DS maintenance testing.
+
+Polars supports all maintenance operations via read-modify-write pattern:
+- INSERT: Append new data to Parquet files
+- DELETE: Read-filter-write pattern (rewrites table)
+- UPDATE: Read-modify-write pattern (rewrites table)
+- MERGE: Read-join-write pattern (rewrites table)
+
+Implementation Note:
+    Polars operates entirely in RAM, so the full table rewrite approach is
+    acceptable and matches Polars' standard data processing patterns. For
+    very large datasets that exceed available RAM, use Delta Lake or Iceberg
+    which support incremental operations with transaction logs.
+
+    The read-modify-write pattern provides:
+    - Full TPC-H and TPC-DS maintenance compliance
+    - Atomic operations via backup-and-swap
+    - Consistent behavior with Polars' immutable data model
+
+Copyright 2026 Joe Harris / BenchBox Project
+
+Licensed under the MIT License. See LICENSE file in the project root for details.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+try:
+    import polars as pl
+
+    POLARS_AVAILABLE = True
+except ImportError:
+    pl = None  # type: ignore[assignment]
+    POLARS_AVAILABLE = False
+
+from benchbox.core.dataframe.maintenance_interface import (
+    POLARS_CAPABILITIES,
+    BaseDataFrameMaintenanceOperations,
+    DataFrameMaintenanceCapabilities,
+)
+
+if TYPE_CHECKING:
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class PolarsMaintenanceOperations(BaseDataFrameMaintenanceOperations):
+    """Polars maintenance operations implementation.
+
+    Implements all TPC maintenance operations via read-modify-write patterns:
+    - INSERT: Append new Parquet files
+    - DELETE: Read-filter-write (rewrites table without matching rows)
+    - UPDATE: Read-modify-write (rewrites table with modified values)
+    - MERGE: Read-join-write (rewrites table with merged data)
+
+    All operations use atomic backup-and-swap to ensure data integrity.
+
+    Note:
+        Polars operates entirely in RAM, so the full table rewrite approach
+        is standard. For datasets exceeding available RAM, use Delta Lake
+        or Iceberg which support incremental operations.
+
+    Example:
+        ops = PolarsMaintenanceOperations()
+
+        # Insert new rows
+        result = ops.insert_rows(
+            table_path="/data/orders",
+            dataframe=new_orders_df,
+            mode="append"
+        )
+
+        # Delete old rows (rewrites table)
+        result = ops.delete_rows(
+            table_path="/data/orders",
+            condition="order_date < '2020-01-01'"
+        )
+
+        # Update rows (rewrites table)
+        result = ops.update_rows(
+            table_path="/data/orders",
+            condition="status = 'pending'",
+            updates={"status": "'cancelled'"}
+        )
+
+        # Merge/upsert (rewrites table)
+        result = ops.merge_rows(
+            table_path="/data/orders",
+            source_dataframe=new_orders,
+            merge_condition="target.id = source.id",
+            when_matched={"status": "source.status"},
+            when_not_matched={"id": "source.id", "status": "source.status"}
+        )
+    """
+
+    def __init__(self, working_dir: str | Path | None = None) -> None:
+        """Initialize Polars maintenance operations.
+
+        Args:
+            working_dir: Optional working directory for temporary files
+
+        Raises:
+            ImportError: If Polars is not installed
+        """
+        super().__init__()
+
+        if not POLARS_AVAILABLE:
+            raise ImportError(
+                "Polars is not installed. Install with: pip install polars\n"
+                "For TPC-H/TPC-DS maintenance tests, install the full package:\n"
+                "pip install 'benchbox[dataframe]'"
+            )
+
+        self.working_dir = Path(working_dir) if working_dir else None
+        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
+
+    def _get_capabilities(self) -> DataFrameMaintenanceCapabilities:
+        """Return Polars maintenance capabilities.
+
+        Returns:
+            POLARS_CAPABILITIES (insert + partition delete only)
+        """
+        return POLARS_CAPABILITIES
+
+    def _do_insert(
+        self,
+        table_path: Path | str,
+        dataframe: Any,
+        partition_columns: list[str] | None,
+        mode: str,
+    ) -> int:
+        """Insert rows by writing to Parquet files.
+
+        Args:
+            table_path: Path to the table directory
+            dataframe: Polars DataFrame to insert
+            partition_columns: Columns to partition by (creates subdirectories)
+            mode: Write mode ("append" or "overwrite")
+
+        Returns:
+            Number of rows inserted
+
+        Raises:
+            ValueError: If mode is invalid
+            TypeError: If dataframe is not a Polars DataFrame
+        """
+        table_path = Path(table_path)
+
+        # Validate and convert dataframe
+        if isinstance(dataframe, pl.LazyFrame):
+            df = dataframe.collect()
+        elif isinstance(dataframe, pl.DataFrame):
+            df = dataframe
+        else:
+            raise TypeError(f"Expected Polars DataFrame or LazyFrame, got {type(dataframe)}")
+
+        row_count = df.height
+
+        if row_count == 0:
+            self.logger.info("No rows to insert")
+            return 0
+
+        # Ensure directory exists
+        table_path.mkdir(parents=True, exist_ok=True)
+
+        if mode == "overwrite":
+            # Remove existing files
+            for f in table_path.glob("*.parquet"):
+                f.unlink()
+
+        if partition_columns:
+            # Write partitioned data
+            self._write_partitioned(df, table_path, partition_columns, mode)
+        else:
+            # Write to a single file or append
+            self._write_single(df, table_path, mode)
+
+        self.logger.info(f"Inserted {row_count} rows to {table_path}")
+        return row_count
+
+    def _write_single(self, df: Any, table_path: Path, mode: str) -> None:
+        """Write DataFrame to a single Parquet file.
+
+        Args:
+            df: Polars DataFrame
+            table_path: Directory path
+            mode: "append" or "overwrite"
+        """
+        if mode == "append":
+            # Generate unique filename
+            existing = list(table_path.glob("part-*.parquet"))
+            part_num = len(existing)
+            file_path = table_path / f"part-{part_num:05d}.parquet"
+        else:
+            file_path = table_path / "part-00000.parquet"
+
+        df.write_parquet(file_path)
+        self.logger.debug(f"Wrote {df.height} rows to {file_path}")
+
+    def _write_partitioned(
+        self,
+        df: Any,
+        table_path: Path,
+        partition_columns: list[str],
+        mode: str,
+    ) -> None:
+        """Write DataFrame with Hive-style partitioning.
+
+        Args:
+            df: Polars DataFrame
+            table_path: Base directory path
+            partition_columns: Columns to partition by
+            mode: Write mode
+        """
+        # Group by partition columns
+        for partition_vals, partition_df in df.group_by(partition_columns):
+            # Build partition path
+            if isinstance(partition_vals, tuple):
+                parts = zip(partition_columns, partition_vals)
+            else:
+                parts = [(partition_columns[0], partition_vals)]
+
+            partition_path = table_path
+            for col, val in parts:
+                partition_path = partition_path / f"{col}={val}"
+
+            partition_path.mkdir(parents=True, exist_ok=True)
+
+            # Write partition data
+            self._write_single(partition_df, partition_path, mode)
+
+    def _do_delete(
+        self,
+        table_path: Path | str,
+        condition: str | Any,
+    ) -> int:
+        """Delete rows using read-filter-write pattern.
+
+        Since Polars doesn't support row-level deletes, we:
+        1. Read all existing data
+        2. Filter out rows matching the condition
+        3. Write the remaining data back
+
+        This is expensive for large datasets. For better performance,
+        use Delta Lake or Iceberg which support row-level deletes.
+
+        Args:
+            table_path: Path to the table directory
+            condition: SQL-like condition string (e.g., "id > 100")
+
+        Returns:
+            Number of rows deleted
+
+        Note:
+            The condition is parsed using Polars SQL syntax.
+            Example conditions:
+            - "order_date < '2020-01-01'"
+            - "status = 'cancelled'"
+            - "amount > 1000 AND region = 'US'"
+        """
+        table_path = Path(table_path)
+
+        if not table_path.exists():
+            self.logger.warning(f"Table path does not exist: {table_path}")
+            return 0
+
+        # Read all existing data
+        parquet_files = list(table_path.glob("**/*.parquet"))
+        if not parquet_files:
+            self.logger.warning(f"No Parquet files found in {table_path}")
+            return 0
+
+        # Read all data
+        self.logger.debug(f"Reading {len(parquet_files)} files from {table_path}")
+        df = pl.scan_parquet(parquet_files).collect()
+        original_count = df.height
+
+        if original_count == 0:
+            return 0
+
+        # Apply filter to keep rows NOT matching the delete condition
+        # We need to negate the condition
+        try:
+            # Use Polars SQL to evaluate the condition
+            df_with_ctx = pl.SQLContext(register_globals=True)
+            df_with_ctx.register("__table__", df)
+
+            # Query to find rows TO DELETE (matching condition)
+            delete_query = f"SELECT COUNT(*) as cnt FROM __table__ WHERE {condition}"
+            delete_count_df = df_with_ctx.execute(delete_query).collect()
+            delete_count = delete_count_df["cnt"][0]
+
+            if delete_count == 0:
+                self.logger.info("No rows match delete condition")
+                return 0
+
+            # Query to keep rows NOT matching condition
+            keep_query = f"SELECT * FROM __table__ WHERE NOT ({condition})"
+            remaining_df = df_with_ctx.execute(keep_query).collect()
+
+        except Exception as e:
+            # Fallback: try to parse condition as Polars expression
+            self.logger.warning(f"SQL condition parsing failed: {e}. Trying expression parse.")
+            raise ValueError(
+                f"Failed to parse delete condition: {condition}\n"
+                f'Use SQL-like syntax: "column > value" or "column = \'value\'"\n'
+                f"Error: {e}"
+            ) from e
+
+        rows_deleted = original_count - remaining_df.height
+
+        if rows_deleted > 0:
+            # Backup and rewrite
+            backup_path = table_path.parent / f"{table_path.name}_backup_{os.getpid()}"
+            try:
+                # Move existing files to backup
+                shutil.move(str(table_path), str(backup_path))
+
+                # Write remaining data
+                table_path.mkdir(parents=True, exist_ok=True)
+                if remaining_df.height > 0:
+                    remaining_df.write_parquet(table_path / "part-00000.parquet")
+
+                # Remove backup on success
+                shutil.rmtree(backup_path)
+
+            except Exception as e:
+                # Restore from backup on failure
+                if backup_path.exists():
+                    shutil.rmtree(table_path, ignore_errors=True)
+                    shutil.move(str(backup_path), str(table_path))
+                raise RuntimeError(f"Delete failed, restored from backup: {e}") from e
+
+        self.logger.info(f"Deleted {rows_deleted} rows from {table_path}")
+        return rows_deleted
+
+    def _do_update(
+        self,
+        table_path: Path | str,
+        condition: str | Any,
+        updates: dict[str, Any],
+    ) -> int:
+        """Update rows using read-modify-write pattern.
+
+        Since Polars uses immutable DataFrames, we:
+        1. Read all existing data
+        2. Apply updates to rows matching the condition
+        3. Write the modified data back
+
+        This is Polars' standard approach for data modification. For datasets
+        exceeding available RAM, use Delta Lake or Iceberg.
+
+        Args:
+            table_path: Path to the table directory
+            condition: SQL-like condition string (e.g., "status = 'pending'")
+            updates: Dict mapping column names to new values or expressions
+
+        Returns:
+            Number of rows updated
+
+        Example:
+            result = ops.update_rows(
+                table_path="/data/orders",
+                condition="status = 'pending'",
+                updates={"status": "'cancelled'", "updated_at": "CURRENT_DATE"}
+            )
+        """
+        table_path = Path(table_path)
+
+        if not table_path.exists():
+            self.logger.warning(f"Table path does not exist: {table_path}")
+            return 0
+
+        # Read all existing data
+        parquet_files = list(table_path.glob("**/*.parquet"))
+        if not parquet_files:
+            self.logger.warning(f"No Parquet files found in {table_path}")
+            return 0
+
+        df = pl.scan_parquet(parquet_files).collect()
+        original_count = df.height
+
+        if original_count == 0:
+            return 0
+
+        try:
+            # Use Polars SQL to identify rows to update and apply changes
+            df_with_ctx = pl.SQLContext(register_globals=True)
+            df_with_ctx.register("__table__", df)
+
+            # Count matching rows
+            count_query = f"SELECT COUNT(*) as cnt FROM __table__ WHERE {condition}"
+            match_count = df_with_ctx.execute(count_query).collect()["cnt"][0]
+
+            if match_count == 0:
+                self.logger.info("No rows match update condition")
+                return 0
+
+            # Polars SQL doesn't support UPDATE directly, so we use CASE expressions
+            # Build a SELECT that applies updates via CASE WHEN
+            select_cols = []
+            for col_name in df.columns:
+                if col_name in updates:
+                    value = updates[col_name]
+                    select_cols.append(f"CASE WHEN ({condition}) THEN {value} ELSE {col_name} END AS {col_name}")
+                else:
+                    select_cols.append(col_name)
+
+            update_query = f"SELECT {', '.join(select_cols)} FROM __table__"
+            updated_df = df_with_ctx.execute(update_query).collect()
+
+        except Exception as e:
+            raise ValueError(
+                f"Failed to parse update condition or values: {condition}\n"
+                f"Updates: {updates}\n"
+                f"Use SQL-like syntax for conditions and values.\n"
+                f"Error: {e}"
+            ) from e
+
+        # Atomic write with backup
+        self._atomic_rewrite(table_path, updated_df)
+
+        self.logger.info(f"Updated {match_count} rows in {table_path}")
+        return match_count
+
+    def _do_merge(
+        self,
+        table_path: Path | str,
+        source_dataframe: Any,
+        merge_condition: str | Any,
+        when_matched: dict[str, Any] | None,
+        when_not_matched: dict[str, Any] | None,
+    ) -> int:
+        """Merge (upsert) rows using read-join-write pattern.
+
+        Since Polars uses immutable DataFrames, we:
+        1. Read target table
+        2. Join with source on merge key
+        3. Apply when_matched updates to matching rows
+        4. Insert when_not_matched rows
+        5. Write the merged data back
+
+        This is Polars' standard approach for upsert operations.
+
+        Args:
+            table_path: Path to the table directory
+            source_dataframe: Polars DataFrame with source data
+            merge_condition: Join condition (e.g., "target.id = source.id")
+            when_matched: Dict of column updates for matched rows
+            when_not_matched: Dict of column values for inserted rows
+
+        Returns:
+            Number of rows affected (updated + inserted)
+
+        Example:
+            result = ops.merge_rows(
+                table_path="/data/dim_customer",
+                source_dataframe=updated_customers,
+                merge_condition="target.c_custkey = source.c_custkey",
+                when_matched={"c_name": "source.c_name", "c_address": "source.c_address"},
+                when_not_matched={"c_custkey": "source.c_custkey", "c_name": "source.c_name"}
+            )
+        """
+        table_path = Path(table_path)
+
+        # Convert source to Polars DataFrame
+        if isinstance(source_dataframe, pl.LazyFrame):
+            source_df = source_dataframe.collect()
+        elif isinstance(source_dataframe, pl.DataFrame):
+            source_df = source_dataframe
+        else:
+            raise TypeError(f"Expected Polars DataFrame or LazyFrame, got {type(source_dataframe)}")
+
+        source_count = source_df.height
+        if source_count == 0:
+            self.logger.info("Source DataFrame is empty, nothing to merge")
+            return 0
+
+        # Read target table
+        if not table_path.exists():
+            # No target table - just insert all source rows
+            self.logger.info(f"Target table doesn't exist, inserting {source_count} rows")
+            return self._do_insert(table_path, source_df, None, "append")
+
+        parquet_files = list(table_path.glob("**/*.parquet"))
+        if not parquet_files:
+            return self._do_insert(table_path, source_df, None, "append")
+
+        target_df = pl.scan_parquet(parquet_files).collect()
+
+        # Parse merge condition to extract key column(s)
+        # Expected format: "target.col = source.col" or just "col"
+        merge_key = self._parse_merge_key(str(merge_condition))
+
+        rows_updated = 0
+        rows_inserted = 0
+
+        # Identify matching and non-matching rows
+        source_keys = set(source_df[merge_key].to_list())
+        target_keys = set(target_df[merge_key].to_list())
+
+        matching_keys = source_keys & target_keys
+        new_keys = source_keys - target_keys
+
+        # Apply when_matched updates
+        if when_matched and matching_keys:
+            # Update matching rows in target with values from source
+            for key_val in matching_keys:
+                # Get source row for this key
+                source_row = source_df.filter(pl.col(merge_key) == key_val)
+
+                # Build update expressions
+                for col, expr in when_matched.items():
+                    if isinstance(expr, str) and expr.startswith("source."):
+                        source_col = expr[7:]  # Remove "source." prefix
+                        new_value = source_row[source_col][0]
+                        target_df = target_df.with_columns(
+                            pl.when(pl.col(merge_key) == key_val)
+                            .then(pl.lit(new_value))
+                            .otherwise(pl.col(col))
+                            .alias(col)
+                        )
+                rows_updated += 1
+
+        # Insert when_not_matched rows
+        if when_not_matched and new_keys:
+            new_rows_df = source_df.filter(pl.col(merge_key).is_in(list(new_keys)))
+            rows_inserted = new_rows_df.height
+
+            # Ensure columns match target schema
+            for col in target_df.columns:
+                if col not in new_rows_df.columns:
+                    # Add missing column with null
+                    new_rows_df = new_rows_df.with_columns(pl.lit(None).alias(col))
+
+            # Select only target columns in correct order
+            new_rows_df = new_rows_df.select(target_df.columns)
+
+            # Append new rows
+            target_df = pl.concat([target_df, new_rows_df])
+
+        # Atomic write
+        self._atomic_rewrite(table_path, target_df)
+
+        total_affected = rows_updated + rows_inserted
+        self.logger.info(f"Merged {total_affected} rows ({rows_updated} updated, {rows_inserted} inserted)")
+        return total_affected
+
+    def _parse_merge_key(self, merge_condition: str) -> str:
+        """Extract merge key column from condition string.
+
+        Args:
+            merge_condition: Condition like "target.id = source.id" or "id"
+
+        Returns:
+            Column name to use as merge key
+        """
+        condition = merge_condition.strip()
+
+        # Handle "target.col = source.col" format
+        if "=" in condition:
+            left, right = condition.split("=", 1)
+            left = left.strip()
+            right = right.strip()
+
+            # Extract column name from "target.col" or "source.col"
+            if "." in left:
+                return left.split(".", 1)[1]
+            return left
+
+        # Handle simple column name
+        return condition
+
+    def _atomic_rewrite(self, table_path: Path, df: Any) -> None:
+        """Atomically rewrite table using backup-and-swap.
+
+        Args:
+            table_path: Path to table directory
+            df: New DataFrame to write
+
+        Raises:
+            RuntimeError: If write fails (original data is restored)
+        """
+        backup_path = table_path.parent / f"{table_path.name}_backup_{os.getpid()}"
+
+        try:
+            # Move existing files to backup
+            if table_path.exists():
+                shutil.move(str(table_path), str(backup_path))
+
+            # Write new data
+            table_path.mkdir(parents=True, exist_ok=True)
+            if df.height > 0:
+                df.write_parquet(table_path / "part-00000.parquet")
+
+            # Remove backup on success
+            if backup_path.exists():
+                shutil.rmtree(backup_path)
+
+        except Exception as e:
+            # Restore from backup on failure
+            if backup_path.exists():
+                shutil.rmtree(table_path, ignore_errors=True)
+                shutil.move(str(backup_path), str(table_path))
+            raise RuntimeError(f"Atomic rewrite failed, restored from backup: {e}") from e
+
+
+def get_polars_maintenance_operations(working_dir: str | Path | None = None) -> PolarsMaintenanceOperations | None:
+    """Get Polars maintenance operations if Polars is available.
+
+    Args:
+        working_dir: Optional working directory
+
+    Returns:
+        PolarsMaintenanceOperations if Polars is available, None otherwise
+    """
+    if not POLARS_AVAILABLE:
+        logger.debug("Polars not available, cannot create maintenance operations")
+        return None
+
+    return PolarsMaintenanceOperations(working_dir=working_dir)

benchbox/platforms/dataframe/pyspark_df.py

@@ -559,6 +559,21 @@ class PySparkDataFrameAdapter(ExpressionFamilyAdapter[PySparkDF, PySparkLazyDF,
         # Return first column value
         return first_row[0]
 
+    def scalar_to_df(self, data: dict[str, Any]) -> PySparkDF:
+        """Create a single-row PySpark DataFrame from scalar values.
+
+        Args:
+            data: Dictionary mapping column names to scalar values
+
+        Returns:
+            PySpark DataFrame with a single row
+        """
+        self._ensure_spark()
+        from pyspark.sql import Row
+
+        row = Row(**data)
+        return self._spark.createDataFrame([row])
+
     # =========================================================================
     # Window Functions
     # =========================================================================