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

benchbox/core/metadata_primitives/benchmark.py

@@ -1173,6 +1173,343 @@ class MetadataPrimitivesBenchmark(BaseBenchmark):
         else:
             raise TypeError(f"Unsupported connection type: {type(connection)}")
 
+    # =========================================================================
+    # DataFrame Support Methods
+    # =========================================================================
+
+    def supports_dataframe_mode(self) -> bool:
+        """Check if this benchmark supports DataFrame execution mode.
+
+        Returns:
+            True - Metadata Primitives supports DataFrame mode via the
+            DataFrameMetadataOperationsManager for schema introspection,
+            catalog operations, and lakehouse metadata.
+        """
+        return True
+
+    def get_dataframe_operations(
+        self,
+        platform_name: str,
+        spark_session: Any = None,
+    ) -> Any:
+        """Get DataFrame metadata operations manager for a platform.
+
+        Returns a DataFrameMetadataOperationsManager configured for the given
+        platform. Use this to execute metadata introspection operations on
+        DataFrame platforms.
+
+        Args:
+            platform_name: Platform name (e.g., "polars-df", "pyspark-df")
+            spark_session: SparkSession instance (required for pyspark-df catalog ops)
+
+        Returns:
+            DataFrameMetadataOperationsManager instance
+
+        Raises:
+            ValueError: If platform doesn't support DataFrame operations
+
+        Example:
+            >>> benchmark = MetadataPrimitivesBenchmark()
+            >>> ops = benchmark.get_dataframe_operations("polars-df")
+            >>> result = ops.execute_get_schema(df)
+        """
+        from benchbox.core.metadata_primitives.dataframe_operations import (
+            get_dataframe_metadata_manager,
+        )
+
+        manager = get_dataframe_metadata_manager(platform_name, spark_session=spark_session)
+        if manager is None:
+            raise ValueError(
+                f"Platform '{platform_name}' does not support DataFrame metadata operations. "
+                f"Supported platforms: polars-df, pandas-df, pyspark-df, datafusion-df"
+            )
+        return manager
+
+    def get_dataframe_capabilities(
+        self,
+        platform_name: str,
+        spark_session: Any = None,
+    ) -> Any:
+        """Get DataFrame metadata capabilities for a platform.
+
+        Returns the DataFrameMetadataCapabilities describing what metadata
+        operations are supported on the given platform.
+
+        Args:
+            platform_name: Platform name (e.g., "polars-df", "pyspark-df")
+            spark_session: SparkSession instance (required for pyspark-df capability detection)
+
+        Returns:
+            DataFrameMetadataCapabilities instance
+
+        Example:
+            >>> benchmark = MetadataPrimitivesBenchmark()
+            >>> caps = benchmark.get_dataframe_capabilities("polars-df")
+            >>> if caps.supports_catalog:
+            ...     # Run catalog operations
+        """
+        from benchbox.core.metadata_primitives.dataframe_operations import (
+            get_dataframe_metadata_manager,
+        )
+
+        manager = get_dataframe_metadata_manager(platform_name, spark_session=spark_session)
+        if manager is None:
+            from benchbox.core.metadata_primitives.dataframe_operations import (
+                get_platform_capabilities,
+            )
+
+            return get_platform_capabilities(platform_name)
+        return manager.get_capabilities()
+
+    def run_dataframe_benchmark(
+        self,
+        platform_name: str,
+        dataframes: dict[str, Any],
+        spark_session: Any = None,
+        categories: list[str] | None = None,
+        iterations: int = 1,
+    ) -> MetadataBenchmarkResult:
+        """Run Metadata Primitives benchmark on DataFrame platforms.
+
+        Executes metadata introspection operations on the provided DataFrames
+        and returns timing results.
+
+        Args:
+            platform_name: Platform name (e.g., "polars-df", "pyspark-df")
+            dataframes: Dictionary mapping table names to DataFrames
+            spark_session: SparkSession instance (required for pyspark-df catalog ops)
+            categories: Optional list of operation categories to run
+                       ("schema", "catalog", "lakehouse", "complexity")
+            iterations: Number of times to run each operation (default: 1)
+
+        Returns:
+            MetadataBenchmarkResult with all operation timings and summary
+
+        Example:
+            >>> import polars as pl
+            >>> benchmark = MetadataPrimitivesBenchmark()
+            >>> dfs = {"orders": pl.read_parquet("orders.parquet")}
+            >>> result = benchmark.run_dataframe_benchmark("polars-df", dfs)
+        """
+        from benchbox.core.metadata_primitives.dataframe_operations import (
+            OPERATION_CATEGORIES,
+            MetadataOperationCategory,
+            MetadataOperationType,
+        )
+
+        manager = self.get_dataframe_operations(platform_name, spark_session=spark_session)
+        capabilities = manager.get_capabilities()
+
+        # Determine which operations to run
+        operations_to_run: list[MetadataOperationType] = []
+
+        if categories:
+            # Filter by category
+            category_set = {MetadataOperationCategory(c.lower()) for c in categories}
+            for op in MetadataOperationType:
+                if OPERATION_CATEGORIES.get(op) in category_set:
+                    if capabilities.supports_operation(op):
+                        operations_to_run.append(op)
+        else:
+            # Run all supported operations
+            operations_to_run = capabilities.get_supported_operations()
+
+        result = MetadataBenchmarkResult()
+        all_results: list[MetadataQueryResult] = []
+
+        # Schema introspection operations run on each DataFrame
+        schema_ops = {
+            MetadataOperationType.LIST_COLUMNS,
+            MetadataOperationType.GET_DTYPES,
+            MetadataOperationType.GET_SCHEMA,
+            MetadataOperationType.DESCRIBE_STATS,
+            MetadataOperationType.ROW_COUNT,
+            MetadataOperationType.COLUMN_COUNT,
+        }
+
+        for _ in range(iterations):
+            for op in operations_to_run:
+                if op in schema_ops:
+                    # Run schema ops on each DataFrame
+                    for table_name, df in dataframes.items():
+                        df_result = self._execute_dataframe_operation(manager, op, df, table_name)
+                        query_result = self._convert_df_result_to_query_result(df_result, table_name)
+                        all_results.append(query_result)
+
+                        if query_result.success:
+                            result.successful_queries += 1
+                        else:
+                            result.failed_queries += 1
+                        result.total_time_ms += query_result.execution_time_ms
+
+                elif op in {
+                    MetadataOperationType.LIST_DATABASES,
+                    MetadataOperationType.LIST_TABLES,
+                }:
+                    # Catalog ops don't need a DataFrame
+                    df_result = self._execute_catalog_operation(manager, op)
+                    query_result = self._convert_df_result_to_query_result(df_result, "catalog")
+                    all_results.append(query_result)
+
+                    if query_result.success:
+                        result.successful_queries += 1
+                    else:
+                        result.failed_queries += 1
+                    result.total_time_ms += query_result.execution_time_ms
+
+        result.total_queries = len(all_results)
+        result.results = all_results
+
+        # Build category summary
+        self._build_dataframe_category_summary(result, all_results)
+
+        return result
+
+    def _execute_dataframe_operation(
+        self,
+        manager: Any,
+        operation: Any,
+        dataframe: Any,
+        table_name: str,
+    ) -> Any:
+        """Execute a single DataFrame metadata operation.
+
+        Args:
+            manager: DataFrameMetadataOperationsManager
+            operation: MetadataOperationType
+            dataframe: The DataFrame to introspect
+            table_name: Name of the table (for logging)
+
+        Returns:
+            DataFrameMetadataResult
+        """
+        from benchbox.core.metadata_primitives.dataframe_operations import (
+            MetadataOperationType,
+        )
+
+        op_map = {
+            MetadataOperationType.LIST_COLUMNS: manager.execute_list_columns,
+            MetadataOperationType.GET_DTYPES: manager.execute_get_dtypes,
+            MetadataOperationType.GET_SCHEMA: manager.execute_get_schema,
+            MetadataOperationType.DESCRIBE_STATS: manager.execute_describe_stats,
+            MetadataOperationType.ROW_COUNT: manager.execute_row_count,
+            MetadataOperationType.COLUMN_COUNT: manager.execute_column_count,
+        }
+
+        if operation in op_map:
+            return op_map[operation](dataframe)
+
+        # Unsupported operation
+        from benchbox.core.metadata_primitives.dataframe_operations import (
+            DataFrameMetadataResult,
+        )
+
+        return DataFrameMetadataResult.failure_result(
+            operation,
+            f"Operation {operation.value} not supported for DataFrame introspection",
+        )
+
+    def _execute_catalog_operation(
+        self,
+        manager: Any,
+        operation: Any,
+    ) -> Any:
+        """Execute a catalog metadata operation.
+
+        Args:
+            manager: DataFrameMetadataOperationsManager
+            operation: MetadataOperationType
+
+        Returns:
+            DataFrameMetadataResult
+        """
+        from benchbox.core.metadata_primitives.dataframe_operations import (
+            MetadataOperationType,
+        )
+
+        if operation == MetadataOperationType.LIST_DATABASES:
+            return manager.execute_list_databases()
+        elif operation == MetadataOperationType.LIST_TABLES:
+            return manager.execute_list_tables()
+
+        from benchbox.core.metadata_primitives.dataframe_operations import (
+            DataFrameMetadataResult,
+        )
+
+        return DataFrameMetadataResult.failure_result(
+            operation,
+            f"Catalog operation {operation.value} not implemented",
+        )
+
+    def _convert_df_result_to_query_result(
+        self,
+        df_result: Any,
+        table_name: str,
+    ) -> MetadataQueryResult:
+        """Convert DataFrameMetadataResult to MetadataQueryResult.
+
+        Args:
+            df_result: DataFrameMetadataResult
+            table_name: Name of the table
+
+        Returns:
+            MetadataQueryResult
+        """
+        from benchbox.core.metadata_primitives.dataframe_operations import (
+            OPERATION_CATEGORIES,
+        )
+
+        category = OPERATION_CATEGORIES.get(df_result.operation_type)
+        category_str = category.value if category else "schema"
+
+        return MetadataQueryResult(
+            query_id=f"df_{df_result.operation_type.value}_{table_name}",
+            category=category_str,
+            execution_time_ms=df_result.duration_ms,
+            row_count=df_result.result_count,
+            success=df_result.success,
+            error=df_result.error_message,
+        )
+
+    def _build_dataframe_category_summary(
+        self,
+        result: MetadataBenchmarkResult,
+        all_results: list[MetadataQueryResult],
+    ) -> None:
+        """Build category summary from DataFrame benchmark results.
+
+        Args:
+            result: MetadataBenchmarkResult to update
+            all_results: List of MetadataQueryResult
+        """
+        category_times: dict[str, list[float]] = {}
+        category_counts: dict[str, int] = {}
+        category_successes: dict[str, int] = {}
+
+        for qr in all_results:
+            cat = qr.category
+            if cat not in category_times:
+                category_times[cat] = []
+                category_counts[cat] = 0
+                category_successes[cat] = 0
+
+            category_times[cat].append(qr.execution_time_ms)
+            category_counts[cat] += 1
+            if qr.success:
+                category_successes[cat] += 1
+
+        for cat in category_times:
+            times = category_times[cat]
+            result.category_summary[cat] = {
+                "total_queries": category_counts[cat],
+                "successful": category_successes[cat],
+                "failed": category_counts[cat] - category_successes[cat],
+                "total_time_ms": sum(times),
+                "avg_time_ms": sum(times) / len(times) if times else 0.0,
+                "min_time_ms": min(times) if times else 0.0,
+                "max_time_ms": max(times) if times else 0.0,
+            }
+
 
 __all__ = [
     "AclBenchmarkResult",