chore: import upstream snapshot with attribution

python/ray/data/stats.py

@@ -0,0 +1,632 @@
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+
+import pandas as pd
+import pyarrow as pa
+
+from ray.data._internal.tensor_extensions.arrow import convert_to_pyarrow_array
+from ray.data.aggregate import (
+    AggregateFnV2,
+    ApproximateQuantile,
+    ApproximateTopK,
+    Count,
+    Max,
+    Mean,
+    Min,
+    MissingValuePercentage,
+    Std,
+    ZeroPercentage,
+)
+from ray.util.annotations import PublicAPI
+
+if TYPE_CHECKING:
+    from ray.data.dataset import Schema
+    from ray.data.datatype import DataType, TypeCategory
+
+
+logger = logging.getLogger(__name__)
+
+
+@PublicAPI(stability="alpha")
+@dataclass
+class DatasetSummary:
+    """Wrapper for dataset summary statistics.
+
+    Provides methods to access computed statistics.
+
+    Attributes:
+        dataset_schema: PyArrow schema of the original dataset
+    """
+
+    STATISTIC_COLUMN = "statistic"
+
+    # PyArrow requires tables whereby each column's value conforms to the column's dtype as defined by the schema.
+    # However, aggregation results might produce statistics with types different from
+    # the original column (e.g., 'count' is int64 even for string columns).
+    # To handle this, we split statistics into two tables:
+    # 1. _stats_matching_column_dtype: Statistics that share the same type as the
+    #    original column (e.g., min/max for numerical columns). These preserve
+    #    the original column's dtype.
+    # 2. _stats_mismatching_column_dtype: Statistics with different types (e.g., count,
+    #    missing_pct). These use inferred types (e.g., float64 for count).
+    _stats_matching_column_dtype: pa.Table
+    _stats_mismatching_column_dtype: pa.Table
+    dataset_schema: pa.Schema
+    columns: list[str]
+
+    def _safe_convert_table(self, table: pa.Table):
+        """Safely convert a PyArrow table to pandas, handling problematic extension types.
+
+        Args:
+            table: PyArrow table to convert
+
+        Returns:
+            pandas DataFrame with converted data
+        """
+        from ray.data.block import BlockAccessor
+
+        try:
+            return BlockAccessor.for_block(table).to_pandas()
+        except (TypeError, ValueError, pa.ArrowInvalid) as e:
+            logger.warning(
+                f"Direct conversion to pandas failed ({e}), "
+                "attempting column-by-column conversion"
+            )
+
+            result_data = {}
+            for col_name in table.schema.names:
+                col = table.column(col_name)
+                try:
+                    result_data[col_name] = col.to_pandas()
+                except (TypeError, ValueError, pa.ArrowInvalid):
+                    # Cast problematic columns to null type
+                    null_col = pa.nulls(len(col), type=pa.null())
+                    result_data[col_name] = null_col.to_pandas()
+
+            return pd.DataFrame(result_data)
+
+    def _set_statistic_index(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Set the statistic column as index if it exists, else return empty DataFrame.
+
+        Args:
+            df: DataFrame to set index on
+
+        Returns:
+            DataFrame with statistic column as index, or empty DataFrame if column missing
+        """
+        if self.STATISTIC_COLUMN in df.columns:
+            return df.set_index(self.STATISTIC_COLUMN)
+        return pd.DataFrame()
+
+    def to_pandas(self):
+        """Convert summary to a single pandas DataFrame.
+
+        Combines statistics from both schema-matching and schema-changing tables.
+
+        Note: Some PyArrow extension types (like TensorExtensionType) may fail to convert
+        to pandas when all values in a column are None. In such cases, this method
+        attempts to convert column-by-column, casting problematic columns to null type.
+
+        Returns:
+            DataFrame with all statistics, where rows are unique statistics from both tables
+        """
+        df_matching = self._set_statistic_index(
+            self._safe_convert_table(self._stats_matching_column_dtype)
+        )
+        df_changing = self._set_statistic_index(
+            self._safe_convert_table(self._stats_mismatching_column_dtype)
+        )
+
+        # Handle case where both are empty
+        if df_matching.empty and df_changing.empty:
+            return pd.DataFrame(columns=[self.STATISTIC_COLUMN])
+
+        # Combine tables: prefer schema_matching values, fill with schema_changing
+        result = df_matching.combine_first(df_changing)
+
+        return (
+            result.reset_index()
+            .sort_values(self.STATISTIC_COLUMN)
+            .reset_index(drop=True)
+        )
+
+    def _extract_column_from_table(
+        self, table: pa.Table, column: str
+    ) -> Optional[dict]:
+        """Extract a column from a PyArrow table if it exists.
+
+        Args:
+            table: PyArrow table to extract from
+            column: Column name to extract
+
+        Returns:
+            DataFrame with 'statistic' and 'value' columns, or None if column doesn't exist
+        """
+        if column not in table.schema.names:
+            return None
+
+        df = self._safe_convert_table(table)[[self.STATISTIC_COLUMN, column]]
+        return df.rename(columns={column: "value"})
+
+    def get_column_stats(self, column: str):
+        """Get all statistics for a specific column, merging from both tables.
+
+        Args:
+            column: Column name to get statistics for
+
+        Returns:
+            DataFrame with all statistics for the column
+        """
+        dfs = [
+            df
+            for table in [
+                self._stats_matching_column_dtype,
+                self._stats_mismatching_column_dtype,
+            ]
+            if (df := self._extract_column_from_table(table, column)) is not None
+        ]
+
+        if not dfs:
+            raise ValueError(f"Column '{column}' not found in summary tables")
+
+        # Concatenate and merge duplicate statistics (prefer non-null values)
+        combined = pd.concat(dfs, ignore_index=True)
+
+        # Group by statistic and take first non-null value for each group
+        def first_non_null(series):
+            non_null = series.dropna()
+            return non_null.iloc[0] if len(non_null) > 0 else None
+
+        result = (
+            combined.groupby(self.STATISTIC_COLUMN, sort=False)["value"]
+            .apply(first_non_null)
+            .reset_index()
+            .sort_values(self.STATISTIC_COLUMN)
+            .reset_index(drop=True)
+        )
+
+        return result
+
+
+@dataclass
+class _DtypeAggregators:
+    """Container for columns and their aggregators.
+
+    Attributes:
+        column_to_dtype: Mapping from column name to dtype string representation
+        aggregators: List of all aggregators to apply
+    """
+
+    column_to_dtype: Dict[str, str]
+    aggregators: List[AggregateFnV2]
+
+
+def _numerical_aggregators(column: str) -> List[AggregateFnV2]:
+    """Generate default metrics for numerical columns.
+
+    This function returns a list of aggregators that compute the following metrics:
+    - count
+    - mean
+    - min
+    - max
+    - std
+    - approximate_quantile (median)
+    - missing_value_percentage
+    - zero_percentage
+
+    Args:
+        column: The name of the numerical column to compute metrics for.
+
+    Returns:
+        A list of AggregateFnV2 instances that can be used with Dataset.aggregate()
+    """
+    return [
+        Count(on=column, ignore_nulls=False),
+        Mean(on=column, ignore_nulls=True),
+        Min(on=column, ignore_nulls=True),
+        Max(on=column, ignore_nulls=True),
+        Std(on=column, ignore_nulls=True, ddof=0),
+        ApproximateQuantile(on=column, quantiles=[0.5]),
+        MissingValuePercentage(on=column),
+        ZeroPercentage(on=column, ignore_nulls=True),
+    ]
+
+
+def _temporal_aggregators(column: str) -> List[AggregateFnV2]:
+    """Generate default metrics for temporal columns.
+
+    This function returns a list of aggregators that compute the following metrics:
+    - count
+    - min
+    - max
+    - missing_value_percentage
+
+    Args:
+        column: The name of the temporal column to compute metrics for.
+
+    Returns:
+        A list of AggregateFnV2 instances that can be used with Dataset.aggregate()
+    """
+    return [
+        Count(on=column, ignore_nulls=False),
+        Min(on=column, ignore_nulls=True),
+        Max(on=column, ignore_nulls=True),
+        MissingValuePercentage(on=column),
+    ]
+
+
+def _basic_aggregators(column: str) -> List[AggregateFnV2]:
+    """Generate default metrics for all columns.
+
+    This function returns a list of aggregators that compute the following metrics:
+    - count
+    - missing_value_percentage
+    - approximate_top_k (top 10 most frequent values)
+
+    Args:
+        column: The name of the column to compute metrics for.
+
+    Returns:
+        A list of AggregateFnV2 instances that can be used with Dataset.aggregate()
+    """
+    return [
+        Count(on=column, ignore_nulls=False),
+        MissingValuePercentage(on=column),
+        ApproximateTopK(on=column, k=10),
+    ]
+
+
+def _default_dtype_aggregators() -> Dict[
+    Union["DataType", "TypeCategory"], Callable[[str], List[AggregateFnV2]]
+]:
+    """Get default mapping from Ray Data DataType to aggregator factory functions.
+
+    This function returns factory functions that create aggregators for specific columns.
+
+    Returns:
+        Dict mapping DataType or TypeCategory to factory functions that take a column name
+        and return a list of aggregators for that column.
+
+    Examples:
+        >>> from ray.data.datatype import DataType
+        >>> from ray.data.stats import _default_dtype_aggregators
+        >>> mapping = _default_dtype_aggregators()
+        >>> factory = mapping.get(DataType.int32())
+        >>> aggs = factory("my_column")  # Creates aggregators for "my_column"
+    """
+    from ray.data.datatype import DataType, TypeCategory
+
+    # Use pattern-matching types for cleaner mapping
+    return {
+        # Numerical types
+        DataType.int8(): _numerical_aggregators,
+        DataType.int16(): _numerical_aggregators,
+        DataType.int32(): _numerical_aggregators,
+        DataType.int64(): _numerical_aggregators,
+        DataType.uint8(): _numerical_aggregators,
+        DataType.uint16(): _numerical_aggregators,
+        DataType.uint32(): _numerical_aggregators,
+        DataType.uint64(): _numerical_aggregators,
+        DataType.float32(): _numerical_aggregators,
+        DataType.float64(): _numerical_aggregators,
+        DataType.bool(): _numerical_aggregators,
+        # String and binary types
+        DataType.string(): _basic_aggregators,
+        DataType.binary(): _basic_aggregators,
+        # Temporal types - pattern matches all temporal types (timestamp, date, time, duration)
+        TypeCategory.TEMPORAL: _temporal_aggregators,
+        # Note: Complex types like lists, structs, maps use fallback logic
+        # in _get_aggregators_for_dtype since they can't be easily enumerated
+    }
+
+
+def _get_fallback_aggregators(column: str, dtype: "DataType") -> List[AggregateFnV2]:
+    """Get aggregators using heuristic-based type detection.
+
+    This is a fallback when no explicit mapping is found for the dtype.
+
+    Args:
+        column: Column name
+        dtype: Ray Data DataType for the column
+
+    Returns:
+        List of aggregators suitable for the column type
+    """
+    try:
+        # Check for null type first
+        if dtype.is_arrow_type() and pa.types.is_null(dtype._physical_dtype):
+            return [Count(on=column, ignore_nulls=False)]
+        elif dtype.is_numerical_type():
+            return _numerical_aggregators(column)
+        elif dtype.is_temporal_type():
+            return _temporal_aggregators(column)
+        else:
+            # Default for strings, binary, lists, nested types, etc.
+            return _basic_aggregators(column)
+
+    except Exception as e:
+        logger.warning(
+            f"Could not determine aggregators for column '{column}' with dtype {dtype}: {e}. "
+            f"Using basic aggregators."
+        )
+        return _basic_aggregators(column)
+
+
+def _get_aggregators_for_dtype(
+    column: str,
+    dtype: "DataType",
+    dtype_agg_mapping: Dict[
+        Union["DataType", "TypeCategory"], Callable[[str], List[AggregateFnV2]]
+    ],
+) -> List[AggregateFnV2]:
+    """Get aggregators for a specific column based on its DataType.
+
+    Attempts to match the dtype against the provided mapping first, then
+    falls back to heuristic-based selection if no match is found.
+
+    Args:
+        column: Column name
+        dtype: Ray Data DataType for the column
+        dtype_agg_mapping: Mapping from DataType to factory functions
+
+    Returns:
+        List of aggregators with the column name properly set
+    """
+    from ray.data.datatype import DataType, TypeCategory
+
+    # Try to find a match in the mapping
+    for mapping_key, factory in dtype_agg_mapping.items():
+        if isinstance(mapping_key, DataType) and dtype == mapping_key:
+            return factory(column)
+        elif isinstance(mapping_key, (TypeCategory, str)) and dtype.is_of(mapping_key):
+            return factory(column)
+
+    # Fallback: Use heuristic-based selection
+    return _get_fallback_aggregators(column, dtype)
+
+
+def _dtype_aggregators_for_dataset(
+    schema: Optional["Schema"],
+    columns: Optional[List[str]] = None,
+    dtype_agg_mapping: Optional[
+        Dict[Union["DataType", "TypeCategory"], Callable[[str], List[AggregateFnV2]]]
+    ] = None,
+) -> _DtypeAggregators:
+    """Generate aggregators for columns in a dataset based on their DataTypes.
+
+    Args:
+        schema: A Ray Schema instance
+        columns: List of columns to include. If None, all columns will be included.
+        dtype_agg_mapping: Optional user-provided mapping from DataType to aggregator factories.
+            Each value should be a callable that takes a column name and returns aggregators.
+            This will be merged with the default mapping (user mapping takes precedence).
+
+    Returns:
+        _DtypeAggregators containing column-to-dtype mapping and aggregators
+
+    Raises:
+        ValueError: If schema is None or if specified columns don't exist in schema
+    """
+    from ray.data.datatype import DataType
+
+    if not schema:
+        raise ValueError("Dataset must have a schema to determine column types")
+
+    if columns is None:
+        columns = schema.names
+
+    # Validate columns exist in schema
+    missing_cols = set(columns) - set(schema.names)
+    if missing_cols:
+        raise ValueError(f"Columns {missing_cols} not found in dataset schema")
+
+    # Build final mapping: default + user overrides
+    defaults = _default_dtype_aggregators()
+    if dtype_agg_mapping:
+        # Put user overrides first so they are checked before default patterns
+        final_mapping = dtype_agg_mapping.copy()
+        for k, v in defaults.items():
+            if k not in final_mapping:
+                final_mapping[k] = v
+    else:
+        final_mapping = defaults
+
+    # Generate aggregators for each column
+    column_to_dtype = {}
+    all_aggs = []
+    name_to_type = dict(zip(schema.names, schema.types))
+
+    for name in columns:
+        pa_type = name_to_type[name]
+        if pa_type is None or pa_type is object:
+            logger.warning(f"Skipping field '{name}': type is None or unsupported")
+            continue
+
+        ray_dtype = DataType.from_arrow(pa_type)
+        column_to_dtype[name] = str(ray_dtype)
+        all_aggs.extend(_get_aggregators_for_dtype(name, ray_dtype, final_mapping))
+
+    return _DtypeAggregators(
+        column_to_dtype=column_to_dtype,
+        aggregators=all_aggs,
+    )
+
+
+def _format_stats(
+    agg: AggregateFnV2, value: Any, agg_type: pa.DataType
+) -> Dict[str, Tuple[Any, pa.DataType]]:
+    """Format aggregation result into stat entries.
+
+    Takes the raw aggregation result and formats it into one or more stat
+    entries. For scalar results, returns a single entry. For list results,
+    expands into multiple indexed entries.
+
+    Args:
+        agg: The aggregator instance
+        value: The aggregation result value
+        agg_type: PyArrow type of the aggregation result
+
+    Returns:
+        Dictionary mapping stat names to (value, type) tuples
+    """
+    from ray.data.datatype import DataType
+
+    agg_name = agg.get_agg_name()
+
+    # Handle list results: expand into separate indexed stats
+    # If the value is None but the type is list, it means we got a null result
+    # for a list-type aggregator (e.g., ignore_nulls=True and all nulls).
+    is_list_type = (
+        pa.types.is_list(agg_type) or DataType.from_arrow(agg_type).is_list_type()
+    )
+
+    if isinstance(value, list) or (value is None and is_list_type):
+        scalar_type = (
+            agg_type.value_type
+            if DataType.from_arrow(agg_type).is_list_type()
+            else agg_type
+        )
+        if value is None:
+            # Can't expand None without knowing the size, return as-is
+            pass
+        else:
+            labels = [str(idx) for idx in range(len(value))]
+            return {
+                f"{agg_name}[{label}]": (list_val, scalar_type)
+                for label, list_val in zip(labels, value)
+            }
+
+    # Fallback to scalar result for non-list values or unexpandable Nones
+    return {agg_name: (value, agg_type)}
+
+
+def _parse_summary_stats(
+    agg_result: Dict[str, any],
+    original_schema: pa.Schema,
+    agg_schema: pa.Schema,
+    aggregators: List[AggregateFnV2],
+) -> tuple:
+    """Parse aggregation results into schema-matching and schema-changing stats.
+
+    Args:
+        agg_result: Dictionary of aggregation results with keys like "count(col)"
+        original_schema: Original dataset schema
+        agg_schema: Schema of aggregation results
+        aggregators: List of aggregators used to generate the results
+
+    Returns:
+        Tuple of (schema_matching_stats, schema_changing_stats, column_names)
+    """
+    schema_matching = {}
+    schema_changing = {}
+    columns = set()
+
+    # Build a lookup map from "stat_name(col_name)" to aggregator
+    agg_lookup = {agg.name: agg for agg in aggregators}
+
+    for key, value in agg_result.items():
+        if "(" not in key or not key.endswith(")"):
+            continue
+
+        # Get aggregator and extract info
+        agg = agg_lookup.get(key)
+        if not agg:
+            continue
+
+        col_name = agg.get_target_column()
+        if not col_name:
+            # Skip aggregations without a target column (e.g., Count())
+            continue
+
+        # Format the aggregation results
+        agg_type = agg_schema.field(key).type
+        original_type = original_schema.field(col_name).type
+        formatted_stats = _format_stats(agg, value, agg_type)
+
+        for stat_name, (stat_value, stat_type) in formatted_stats.items():
+            # Add formatted stats to appropriate dict based on schema matching
+            stats_dict = (
+                schema_matching if stat_type == original_type else schema_changing
+            )
+            stats_dict.setdefault(stat_name, {})[col_name] = (stat_value, stat_type)
+
+        columns.add(col_name)
+
+    return schema_matching, schema_changing, columns
+
+
+def _create_pyarrow_array(
+    col_data: List, col_type: Optional[pa.DataType] = None, col_name: str = ""
+) -> pa.Array:
+    """Create a PyArrow array with fallback strategies.
+
+    Uses convert_to_pyarrow_array from arrow_block.py for type inference and
+    error handling when no specific type is provided.
+
+    Args:
+        col_data: List of column values
+        col_type: Optional PyArrow type to use
+        col_name: Column name for error messages (optional)
+
+    Returns:
+        PyArrow array
+    """
+    if col_type is not None:
+        try:
+            return pa.array(col_data, type=col_type)
+        except (pa.ArrowTypeError, pa.ArrowInvalid):
+            # Type mismatch - fall through to type inference
+            pass
+
+    # Use convert_to_pyarrow_array for type inference and error handling
+    # This handles tensors, extension types, and fallback to ArrowPythonObjectArray
+    return convert_to_pyarrow_array(col_data, col_name or "column")
+
+
+def _build_summary_table(
+    stats_dict: Dict[str, Dict[str, tuple]],
+    all_columns: set,
+    original_schema: pa.Schema,
+    preserve_types: bool,
+) -> pa.Table:
+    """Build a PyArrow table from parsed statistics.
+
+    Args:
+        stats_dict: Nested dict of {stat_name: {col_name: (value, type)}}
+        all_columns: Set of all column names across both tables
+        original_schema: Original dataset schema
+        preserve_types: If True, use original schema types for columns
+
+    Returns:
+        PyArrow table with statistics
+    """
+    if not stats_dict:
+        return pa.table({})
+
+    stat_names = sorted(stats_dict.keys())
+    table_data = {DatasetSummary.STATISTIC_COLUMN: stat_names}
+
+    for col_name in sorted(all_columns):
+        # Collect values and infer type
+        col_data = []
+        first_type = None
+
+        for stat_name in stat_names:
+            if col_name in stats_dict[stat_name]:
+                value, agg_type = stats_dict[stat_name][col_name]
+                col_data.append(value)
+                if first_type is None:
+                    first_type = agg_type
+            else:
+                col_data.append(None)
+
+        # Determine column type: prefer original schema, then first aggregation type, then infer
+        if preserve_types and col_name in original_schema.names:
+            col_type = original_schema.field(col_name).type
+        else:
+            col_type = first_type
+
+        table_data[col_name] = _create_pyarrow_array(col_data, col_type, col_name)
+
+    return pa.table(table_data)