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

benchbox/core/transaction_primitives/generator.py

@@ -31,6 +31,7 @@ from benchbox.core.tpch.generator import TPCHDataGenerator
 from benchbox.utils.cloud_storage import CloudStorageGeneratorMixin, create_path_handler
 from benchbox.utils.compression_mixin import CompressionMixin
 from benchbox.utils.datagen_manifest import DataGenerationManifest, resolve_compression_metadata
+from benchbox.utils.file_format import detect_compression
 from benchbox.utils.path_utils import get_benchmark_runs_datagen_path
 from benchbox.utils.verbosity import VerbosityMixin, compute_verbosity
 
@@ -674,7 +675,8 @@ class TransactionPrimitivesDataGenerator(CompressionMixin, CloudStorageGenerator
                 break
 
             # Decompress if needed based on file extension
-            if file_path.suffix == ".zst":
+            compression = detect_compression(file_path)
+            if compression == "zstd":
                 try:
                     import zstandard as zstd
 
@@ -686,14 +688,19 @@ class TransactionPrimitivesDataGenerator(CompressionMixin, CloudStorageGenerator
                 except ImportError:
                     self.log_verbose(f"zstandard not available, skipping {file_path}")
                     continue
-            elif file_path.suffix == ".gz":
+            elif compression == "gzip":
                 with gzip.open(file_path, "rt", encoding="utf-8") as f:
                     rows.extend(self._parse_tbl_stream(f, limit - len(rows)))
-            elif file_path.suffix == ".bz2":
+            elif compression == "bzip2":
                 with bz2.open(file_path, "rt", encoding="utf-8") as f:
                     rows.extend(self._parse_tbl_stream(f, limit - len(rows)))
+            elif compression == "xz":
+                import lzma
+
+                with lzma.open(file_path, "rt", encoding="utf-8") as f:
+                    rows.extend(self._parse_tbl_stream(f, limit - len(rows)))
             else:
-                # Uncompressed .tbl file or numbered shard
+                # Uncompressed .tbl file or numbered shard, or unsupported compression
                 with open(file_path, encoding="utf-8") as f:
                     rows.extend(self._parse_tbl_stream(f, limit - len(rows)))
 

benchbox/core/visualization/__init__.py

@@ -9,7 +9,7 @@ from benchbox.core.visualization.charts import (
     DistributionBoxPlot,
     DistributionSeries,
     PerformanceBarChart,
-    QueryVarianceHeatmap,
+    QueryHeatmap,
     TimeSeriesLineChart,
     TimeSeriesPoint,
 )
@@ -27,7 +27,7 @@ __all__ = [
     "DistributionBoxPlot",
     "DistributionSeries",
     "PerformanceBarChart",
-    "QueryVarianceHeatmap",
+    "QueryHeatmap",
     "TimeSeriesLineChart",
     "TimeSeriesPoint",
     "SUPPORTED_EXPORT_FORMATS",

benchbox/core/visualization/charts.py

@@ -310,8 +310,8 @@ class CostPerformanceScatterPlot(BaseChart):
         )
 
 
-class QueryVarianceHeatmap(BaseChart):
-    """Query × platform variance heatmap."""
+class QueryHeatmap(BaseChart):
+    """Query × platform execution time heatmap."""
 
     def __init__(
         self,
@@ -319,13 +319,13 @@ class QueryVarianceHeatmap(BaseChart):
         queries: Sequence[str],
         platforms: Sequence[str],
         title: str | None = None,
-        colorbar_title: str = "Coefficient of Variation",
+        colorbar_title: str = "Execution Time (ms)",
     ):
         super().__init__()
         self.matrix = matrix
         self.queries = queries
         self.platforms = platforms
-        self.title = title or "Query Variance Heatmap"
+        self.title = title or "Query Execution Time Heatmap"
         self.colorbar_title = colorbar_title
 
     def figure(self):

benchbox/core/visualization/dependencies.py

@@ -15,17 +15,6 @@ def require_plotly() -> tuple[Any, Any]:
     except ImportError as exc:  # pragma: no cover - exercised via dependency check
         raise VisualizationDependencyError(
             "plotly",
-            "Install visualization extras with: "
-            'uv add "plotly>=5.24.0" "kaleido>=0.2.1" "pandas>=2.0.0" "pillow>=10.0.0"',
+            'Install visualization extras with: uv add "plotly>=5.24.0"',
         ) from exc
     return go, pio
-
-
-def require_kaleido() -> Any:
-    """Ensure Kaleido is available for static exports."""
-    _, pio = require_plotly()
-    try:
-        import kaleido  # noqa: F401
-    except ImportError as exc:  # pragma: no cover - exercised via dependency check
-        raise VisualizationDependencyError("kaleido", 'Install with: uv add "kaleido>=0.2.1"') from exc
-    return pio

benchbox/core/visualization/exporters.py

@@ -6,11 +6,11 @@ from collections.abc import Mapping, Sequence
 from pathlib import Path
 from typing import Any
 
-from benchbox.core.visualization.dependencies import require_kaleido, require_plotly
-from benchbox.core.visualization.exceptions import VisualizationDependencyError, VisualizationError
+from benchbox.core.visualization.dependencies import require_plotly
+from benchbox.core.visualization.exceptions import VisualizationError
 from benchbox.core.visualization.styles import ThemeSettings, apply_theme, get_theme
 
-SUPPORTED_EXPORT_FORMATS = {"png", "svg", "pdf", "html"}
+SUPPORTED_EXPORT_FORMATS = {"html"}
 
 
 def export_figure(
@@ -18,9 +18,9 @@ def export_figure(
     output_dir: str | Path,
     base_name: str,
     formats: Sequence[str] | None = None,
-    dpi: int = 300,
     metadata: Mapping[str, Any] | None = None,
     theme: ThemeSettings | None = None,
+    dpi: int = 300,
 ) -> dict[str, Path]:
     """Export a Plotly figure to the requested formats.
 
@@ -28,16 +28,16 @@ def export_figure(
         fig: Plotly Figure to export.
         output_dir: Directory to write files into.
         base_name: Base filename without extension.
-        formats: Iterable of formats (png, svg, pdf, html). Defaults to png+html.
-        dpi: Target DPI for raster outputs (png/pdf).
+        formats: Iterable of formats (html only). Defaults to html.
         metadata: Optional metadata to embed in the figure.
         theme: Theme to apply before export.
+        dpi: Resolution for raster exports (unused until png/pdf support added).
 
     Returns:
         Mapping of format -> output path.
     """
     go, pio = require_plotly()
-    export_formats = list(formats) if formats else ["png", "html"]
+    export_formats = list(formats) if formats else ["html"]
 
     unknown = set(export_formats) - SUPPORTED_EXPORT_FORMATS
     if unknown:
@@ -53,29 +53,18 @@ def export_figure(
         fig.update_layout(meta=dict(metadata))
 
     export_map: dict[str, Path] = {}
-    scale = max(dpi / 72.0, 1.0)
 
     for fmt in export_formats:
         target = output_path / f"{base_name}.{fmt}"
         try:
-            if fmt == "html":
-                pio.write_html(
-                    fig,
-                    target,
-                    include_plotlyjs=True,
-                    full_html=True,
-                    default_width="100%",
-                    default_height="100%",
-                )
-            else:
-                require_kaleido()
-                pio.write_image(fig, target, format=fmt, engine="kaleido", scale=scale)
-        except VisualizationDependencyError:
-            raise
-        except ValueError as exc:
-            if "kaleido" in str(exc).lower():
-                raise VisualizationDependencyError("kaleido", 'Install with: uv add "kaleido>=0.2.1"') from exc
-            raise VisualizationError(f"Failed to export {fmt}: {exc}") from exc
+            pio.write_html(
+                fig,
+                target,
+                include_plotlyjs=True,
+                full_html=True,
+                default_width="100%",
+                default_height="100%",
+            )
         except Exception as exc:  # pragma: no cover - defensive catch
             raise VisualizationError(f"Failed to export {fmt}: {exc}") from exc
 

benchbox/core/visualization/result_plotter.py

@@ -13,6 +13,7 @@ from typing import Any
 import benchbox
 from benchbox.core.results.loader import find_latest_result
 from benchbox.core.results.models import BenchmarkResults
+from benchbox.core.results.normalizer import normalize_result_dict
 from benchbox.core.visualization.charts import (
     BarDatum,
     CostPerformancePoint,
@@ -20,7 +21,7 @@ from benchbox.core.visualization.charts import (
     DistributionBoxPlot,
     DistributionSeries,
     PerformanceBarChart,
-    QueryVarianceHeatmap,
+    QueryHeatmap,
     TimeSeriesLineChart,
     TimeSeriesPoint,
 )
@@ -64,8 +65,38 @@ class ResultPlotter:
         if not results:
             raise VisualizationError("No results provided for visualization.")
         self.results = list(results)
+        self._disambiguate_platforms()
         self.theme = theme if isinstance(theme, ThemeSettings) else get_theme(theme)
 
+    def _disambiguate_platforms(self) -> None:
+        """Append scale factor to platform labels when needed to avoid duplicates.
+
+        Only modifies labels when multiple results share the same platform label
+        but have different scale factors. This keeps labels clean when comparing
+        different platforms at the same scale.
+        """
+        from collections import Counter
+
+        # Count occurrences of each platform label
+        platform_counts = Counter(r.platform for r in self.results)
+
+        # Find platforms with duplicates
+        duplicated_platforms = {p for p, count in platform_counts.items() if count > 1}
+        if not duplicated_platforms:
+            return
+
+        # For each duplicated platform, check if scale factors differ
+        for platform in duplicated_platforms:
+            matching = [r for r in self.results if r.platform == platform]
+            scale_factors = {r.scale_factor for r in matching}
+
+            # Only disambiguate if scale factors actually differ
+            if len(scale_factors) > 1:
+                for result in matching:
+                    sf_label = format_scale_factor(result.scale_factor)
+                    # Mutate the platform label to include scale factor
+                    object.__setattr__(result, "platform", f"{result.platform} SF={sf_label}")
+
     # ------------------------------------------------------------------ Loading
     @classmethod
     def from_sources(
@@ -121,50 +152,39 @@ class ResultPlotter:
 
     @staticmethod
     def _normalize_dict(data: dict[str, Any], source_path: Path | None) -> NormalizedResult:
-        benchmark_block = data.get("benchmark", {}) or {}
-        execution_block = data.get("execution", {}) or {}
-        results_block = data.get("results", {}) or {}
-        queries_block = results_block.get("queries", {}) or {}
-        timing_block = results_block.get("timing", {}) or {}
-        cost_block = data.get("cost_summary") or results_block.get("cost_summary") or {}
-
-        benchmark = benchmark_block.get("name") or benchmark_block.get("id") or "unknown"
-        platform = execution_block.get("platform") or (data.get("platform") or {}).get("name", "unknown")
-        scale = benchmark_block.get("scale_factor") or "unknown"
-        execution_id = execution_block.get("id") or execution_block.get("execution_id")
-        timestamp_raw = execution_block.get("timestamp")
-        timestamp = None
-        if timestamp_raw:
-            try:
-                timestamp = datetime.fromisoformat(timestamp_raw)
-            except ValueError:
-                timestamp = None
-
-        total_ms = timing_block.get("total_ms")
-        avg_ms = timing_block.get("avg_ms")
-        success_rate = queries_block.get("success_rate")
-        cost_total = cost_block.get("total_cost")
-
-        queries: list[NormalizedQuery] = []
-        for query in queries_block.get("details", []) or []:
-            queries.append(
-                NormalizedQuery(
-                    query_id=query.get("id") or query.get("query_id") or "unknown",
-                    execution_time_ms=query.get("execution_time_ms"),
-                    status=query.get("status", "UNKNOWN"),
-                )
+        """Convert raw JSON dict to visualization-specific NormalizedResult.
+
+        Uses the shared normalizer for schema detection and field extraction,
+        then converts to the visualization-specific format with platform label
+        decoration for execution mode differentiation.
+        """
+        normalized = normalize_result_dict(data)
+
+        # Decorate platform label with execution mode for visual differentiation
+        platform = normalized.platform
+        if normalized.execution_mode and normalized.execution_mode.lower() not in ("sql", "datagen", "data_only"):
+            platform = f"{platform} ({normalized.execution_mode})"
+
+        # Convert shared NormalizedQuery to visualization NormalizedQuery
+        queries = [
+            NormalizedQuery(
+                query_id=q.query_id,
+                execution_time_ms=q.execution_time_ms,
+                status=q.status,
             )
+            for q in normalized.queries
+        ]
 
         return NormalizedResult(
-            benchmark=str(benchmark),
-            platform=str(platform),
-            scale_factor=scale,
-            execution_id=execution_id,
-            timestamp=timestamp,
-            total_time_ms=total_ms,
-            avg_time_ms=avg_ms,
-            success_rate=success_rate,
-            cost_total=cost_total,
+            benchmark=normalized.benchmark,
+            platform=platform,
+            scale_factor=normalized.scale_factor,
+            execution_id=normalized.execution_id,
+            timestamp=normalized.timestamp,
+            total_time_ms=normalized.total_time_ms,
+            avg_time_ms=normalized.avg_time_ms,
+            success_rate=normalized.success_rate,
+            cost_total=normalized.cost_total,
             queries=queries,
             source_path=source_path,
             raw=data,
@@ -206,9 +226,16 @@ class ResultPlotter:
 
         cost_summary = getattr(result, "cost_summary", None) or {}
 
+        # Build platform label, including execution mode when not default SQL
+        platform = str(getattr(result, "platform", "unknown"))
+        platform_info = getattr(result, "platform_info", None) or {}
+        execution_mode = platform_info.get("execution_mode")
+        if execution_mode and execution_mode.lower() not in ("sql", "datagen", "data_only"):
+            platform = f"{platform} ({execution_mode})"
+
         return NormalizedResult(
             benchmark=str(getattr(result, "benchmark_name", "unknown")),
-            platform=str(getattr(result, "platform", "unknown")),
+            platform=platform,
             scale_factor=getattr(result, "scale_factor", "unknown"),
             execution_id=getattr(result, "execution_id", None),
             timestamp=timestamp,
@@ -237,7 +264,7 @@ class ResultPlotter:
         if template_name:
             template = get_template(template_name)
 
-        chosen_formats = list(formats or (template.formats if template else ("png", "html")))
+        chosen_formats = list(formats or (template.formats if template else ("html",)))
         chosen_chart_types = list(chart_types or (template.chart_types if template else ()))
         if not chosen_chart_types:
             chosen_chart_types = self._suggest_chart_types() if smart else ["performance_bar"]
@@ -268,9 +295,9 @@ class ResultPlotter:
                 output_dir=output_path,
                 base_name=base_name,
                 formats=chosen_formats,
-                dpi=dpi,
                 metadata=metadata,
                 theme=self.theme,
+                dpi=dpi,
             )
             exports[chart_type] = export_paths
         return exports
@@ -325,7 +352,7 @@ class ResultPlotter:
 
     def _render_query_heatmap(self):
         # Build a matrix of query execution times (ms)
-        query_ids = sorted({q.query_id for result in self.results for q in result.queries})
+        query_ids = sorted({q.query_id for result in self.results for q in result.queries}, key=self._natural_sort_key)
         if not query_ids or len(self.results) < 2:
             logger.info("Heatmap requires at least one query and multiple platforms.")
             return None
@@ -339,13 +366,14 @@ class ResultPlotter:
                 row.append(match.execution_time_ms if match else None)
             matrix.append(row)
 
-        chart = QueryVarianceHeatmap(
+        chart = QueryHeatmap(
             matrix=matrix,
             queries=query_ids,
             platforms=platform_names,
-            title=f"{self._benchmark_label()} Query Variance",
+            title=f"{self._benchmark_label()} Query Execution Times",
+            colorbar_title="Execution Time (ms)",
         )
-        return chart.figure(), f"{self._slug_prefix()}query-variance"
+        return chart.figure(), f"{self._slug_prefix()}query-heatmap"
 
     def _render_cost_scatter(self):
         points: list[CostPerformancePoint] = []
@@ -449,6 +477,19 @@ class ResultPlotter:
         except Exception:
             return str(scale)
 
+    @staticmethod
+    def _natural_sort_key(s: str) -> tuple[float, str]:
+        """Sort key for natural ordering of query IDs (e.g., Q1, Q2, Q10 instead of Q1, Q10, Q2)."""
+        import re
+
+        # Extract leading non-digits and trailing number
+        match = re.match(r"^(\D*)(\d+)(.*)$", s)
+        if match:
+            prefix, num, suffix = match.groups()
+            return (float(num), prefix + suffix)
+        # Fallback: no number found, sort alphabetically at the end
+        return (float("inf"), s)
+
     def group_by(self, field: str) -> dict[str, ResultPlotter]:
         """Split results by a field (platform or benchmark) for batch rendering."""
         if field not in {"platform", "benchmark"}:

benchbox/core/visualization/templates.py

@@ -16,7 +16,7 @@ class ChartTemplate:
     name: str
     description: str
     chart_types: Sequence[str]
-    formats: Sequence[str] = field(default_factory=lambda: ("png", "html"))
+    formats: Sequence[str] = field(default_factory=lambda: ("html",))
     options: dict[str, Any] = field(default_factory=dict)
 
 
@@ -25,31 +25,31 @@ _TEMPLATES: dict[str, ChartTemplate] = {
         name="default",
         description="Baseline set for a single benchmark run (bar + box + heatmap when available).",
         chart_types=("performance_bar", "distribution_box", "query_heatmap"),
-        formats=("png", "html"),
+        formats=("html",),
     ),
     "flagship": ChartTemplate(
         name="flagship",
         description="Eight-platform flagship comparison with heatmap and cost frontier.",
         chart_types=("performance_bar", "query_heatmap", "cost_scatter", "distribution_box"),
-        formats=("png", "svg", "html"),
+        formats=("html",),
     ),
     "head_to_head": ChartTemplate(
         name="head_to_head",
         description="Two-platform comparison with win/loss emphasis.",
         chart_types=("performance_bar", "distribution_box", "query_heatmap"),
-        formats=("png", "html"),
+        formats=("html",),
     ),
     "trends": ChartTemplate(
         name="trends",
         description="Multi-period performance trend lines with regression overlay.",
         chart_types=("time_series", "performance_bar"),
-        formats=("png", "svg", "html"),
+        formats=("html",),
     ),
     "cost_optimization": ChartTemplate(
         name="cost_optimization",
         description="Cost breakdown and price/performance frontier.",
         chart_types=("cost_scatter", "performance_bar"),
-        formats=("png", "svg", "html"),
+        formats=("html",),
     ),
 }
 

benchbox/core/write_primitives/__init__.py

@@ -11,6 +11,13 @@ from benchbox.core.write_primitives.benchmark import (
     OperationResult,
     WritePrimitivesBenchmark,
 )
+from benchbox.core.write_primitives.dataframe_operations import (
+    DataFrameWriteCapabilities,
+    DataFrameWriteOperationsManager,
+    DataFrameWriteResult,
+    WriteOperationType,
+    get_dataframe_write_manager,
+)
 from benchbox.core.write_primitives.generator import WritePrimitivesDataGenerator
 from benchbox.core.write_primitives.operations import WriteOperationsManager
 from benchbox.core.write_primitives.schema import (
@@ -42,6 +49,12 @@ __all__ = [
     "OperationResult",
     "WritePrimitivesDataGenerator",
     "WriteOperationsManager",
+    # DataFrame operations
+    "WriteOperationType",
+    "DataFrameWriteCapabilities",
+    "DataFrameWriteResult",
+    "DataFrameWriteOperationsManager",
+    "get_dataframe_write_manager",
     "TABLES",
     "STAGING_TABLES",
     "INSERT_OPS_LINEITEM",

benchbox/core/write_primitives/benchmark.py

@@ -1135,5 +1135,71 @@ class WritePrimitivesBenchmark(BaseBenchmark, OperationExecutor):
 
         return results
 
+    # ========================================================================
+    # DataFrame Mode Support
+    # ========================================================================
+
+    def supports_dataframe_mode(self) -> bool:
+        """Check if this benchmark supports DataFrame execution mode.
+
+        Write Primitives supports DataFrame mode for platforms like Polars,
+        PySpark, and Pandas that can perform write operations.
+
+        Returns:
+            True - Write Primitives supports DataFrame mode
+        """
+        return True
+
+    def get_dataframe_operations(self, platform_name: str) -> "DataFrameWriteOperationsManager | None":
+        """Get DataFrame write operations manager for a platform.
+
+        Args:
+            platform_name: Platform name (e.g., "polars-df", "pyspark-df")
+
+        Returns:
+            DataFrameWriteOperationsManager if platform supports DataFrame writes,
+            None otherwise.
+
+        Example:
+            benchmark = WritePrimitivesBenchmark()
+            manager = benchmark.get_dataframe_operations("polars-df")
+            if manager:
+                result = manager.execute_insert(
+                    table_path="/data/orders",
+                    dataframe=orders_df,
+                    mode="append"
+                )
+        """
+        from benchbox.core.write_primitives.dataframe_operations import (
+            get_dataframe_write_manager,
+        )
+
+        return get_dataframe_write_manager(platform_name)
+
+    def get_dataframe_capabilities(self, platform_name: str) -> "DataFrameWriteCapabilities | None":
+        """Get DataFrame write capabilities for a platform.
+
+        Args:
+            platform_name: Platform name (e.g., "polars-df", "pyspark-df")
+
+        Returns:
+            DataFrameWriteCapabilities if platform supports DataFrame writes,
+            None otherwise.
+        """
+        manager = self.get_dataframe_operations(platform_name)
+        if manager is None:
+            return None
+        return manager.get_capabilities()
+
+
+# Import for type hints
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from benchbox.core.write_primitives.dataframe_operations import (
+        DataFrameWriteCapabilities,
+        DataFrameWriteOperationsManager,
+    )
+
 
 __all__ = ["WritePrimitivesBenchmark", "OperationResult"]