import pandas as pd import pyarrow as pa import pytest from packaging.version import parse as parse_version import ray from ray.data._internal.util import rows_same from ray.data._internal.utils.arrow_utils import get_pyarrow_version from ray.data.aggregate import ( ApproximateQuantile, ApproximateTopK, Count, Max, Mean, Min, MissingValuePercentage, Std, ZeroPercentage, ) from ray.data.datatype import DataType from ray.data.stats import ( DatasetSummary, _basic_aggregators, _default_dtype_aggregators, _dtype_aggregators_for_dataset, _numerical_aggregators, _temporal_aggregators, ) class TestDtypeAggregatorsForDataset: """Test suite for _dtype_aggregators_for_dataset function.""" @pytest.mark.parametrize( "data,expected_dtypes,expected_agg_count", [ # Numerical columns only ( [{"int_col": 1, "float_col": 1.5}], { "int_col": "DataType(arrow:int64)", "float_col": "DataType(arrow:double)", }, 16, # 2 columns * 8 aggregators each ), # Mixed numerical and string ( [{"num": 1, "str": "a"}], {"num": "DataType(arrow:int64)", "str": "DataType(arrow:string)"}, 11, # 1 numerical * 8 + 1 string * 3 ), # Boolean treated as numerical ( [{"bool_col": True, "int_col": 1}], { "bool_col": "DataType(arrow:bool)", "int_col": "DataType(arrow:int64)", }, 16, # 2 columns * 8 aggregators each ), ], ) def test_column_type_detection(self, data, expected_dtypes, expected_agg_count): """Test that column types are correctly detected and mapped.""" ds = ray.data.from_items(data) result = _dtype_aggregators_for_dataset(ds.schema()) assert (result.column_to_dtype, len(result.aggregators)) == ( expected_dtypes, expected_agg_count, ) def test_column_filtering(self): """Test that only specified columns are included.""" data = [{"col1": 1, "col2": "a", "col3": 1.5}] ds = ray.data.from_items(data) result = _dtype_aggregators_for_dataset(ds.schema(), columns=["col1", "col3"]) assert (set(result.column_to_dtype.keys()), len(result.aggregators)) == ( {"col1", "col3"}, 16, ) def test_empty_columns_list(self): """Test behavior with empty columns list.""" data = [{"col1": 1, "col2": "a"}] ds = ray.data.from_items(data) result = _dtype_aggregators_for_dataset(ds.schema(), columns=[]) assert (len(result.column_to_dtype), len(result.aggregators)) == (0, 0) def test_invalid_columns_raises_error(self): """Test error handling when columns parameter contains non-existent columns.""" data = [{"col1": 1}] ds = ray.data.from_items(data) with pytest.raises(ValueError, match="not found in dataset schema"): _dtype_aggregators_for_dataset(ds.schema(), columns=["nonexistent"]) def test_none_schema_raises_error(self): """Test that None schema raises appropriate error.""" with pytest.raises(ValueError, match="must have a schema"): _dtype_aggregators_for_dataset(None) def test_custom_dtype_mapping(self): """Test that custom dtype mappings override defaults.""" data = [{"int_col": 1}] ds = ray.data.from_items(data) # Override int64 to only use Count and Mean custom_mapping = {DataType.int64(): lambda col: [Count(on=col), Mean(on=col)]} result = _dtype_aggregators_for_dataset( ds.schema(), dtype_agg_mapping=custom_mapping ) assert [type(agg) for agg in result.aggregators] == [Count, Mean] @pytest.mark.skipif( get_pyarrow_version() < parse_version("14.0.0"), reason="Requires pyarrow >= 14.0.0", ) def test_custom_dtype_mapping_pattern_precedence(self): """Test that specific custom mappings take precedence over default patterns.""" import datetime # Use from_arrow to ensure we get exactly timestamp[us] t = pa.table( {"ts": pa.array([datetime.datetime(2024, 1, 1)], type=pa.timestamp("us"))} ) ds = ray.data.from_arrow(t) # Override specific timestamp type to only use Count # Default for temporal is [Count, Min, Max, MissingValuePercentage] ts_dtype = DataType.from_arrow(pa.timestamp("us")) custom_mapping = {ts_dtype: lambda col: [Count(on=col)]} result = _dtype_aggregators_for_dataset( ds.schema(), dtype_agg_mapping=custom_mapping ) # Should only have 1 aggregator if our specific override was used. # If the default DataType.temporal() pattern matched first, we'd get 4 aggregators. assert len(result.aggregators) == 1 assert isinstance(result.aggregators[0], Count) @pytest.mark.skipif( get_pyarrow_version() < parse_version("14.0.0"), reason="Requires pyarrow >= 14.0.0", ) @pytest.mark.parametrize( "pa_type", [ pa.timestamp("us"), # Temporal: count, min, max, missing% pa.date32(), # Temporal pa.time64("us"), # Temporal ], ) def test_temporal_types(self, pa_type): """Test that temporal types get appropriate aggregators.""" table = pa.table({"temporal_col": pa.array([1, 2, 3], type=pa_type)}) ds = ray.data.from_arrow(table) result = _dtype_aggregators_for_dataset(ds.schema()) assert "temporal_col" in result.column_to_dtype assert [type(agg) for agg in result.aggregators] == [ Count, Min, Max, MissingValuePercentage, ] class TestIndividualAggregatorFunctions: """Test suite for individual aggregator generator functions.""" def test_numerical_aggregators(self): """Test _numerical_aggregators function.""" aggs = _numerical_aggregators("test_col") assert len(aggs) == 8 assert all(agg.get_target_column() == "test_col" for agg in aggs) assert [type(agg) for agg in aggs] == [ Count, Mean, Min, Max, Std, ApproximateQuantile, MissingValuePercentage, ZeroPercentage, ] def test_temporal_aggregators(self): """Test _temporal_aggregators function.""" aggs = _temporal_aggregators("test_col") assert len(aggs) == 4 assert all(agg.get_target_column() == "test_col" for agg in aggs) assert [type(agg) for agg in aggs] == [Count, Min, Max, MissingValuePercentage] def test_basic_aggregators(self): """Test _basic_aggregators function.""" aggs = _basic_aggregators("test_col") assert len(aggs) == 3 assert all(agg.get_target_column() == "test_col" for agg in aggs) assert [type(agg) for agg in aggs] == [ Count, MissingValuePercentage, ApproximateTopK, ] class TestDefaultDtypeAggregators: """Test suite for _default_dtype_aggregators function.""" @pytest.mark.skipif( get_pyarrow_version() < parse_version("14.0.0"), reason="Requires pyarrow >= 14.0.0", ) @pytest.mark.parametrize( "dtype_factory,expected_agg_types,uses_pattern_matching", [ ( DataType.int32, [ Count, Mean, Min, Max, Std, ApproximateQuantile, MissingValuePercentage, ZeroPercentage, ], False, ), # Numerical ( DataType.float64, [ Count, Mean, Min, Max, Std, ApproximateQuantile, MissingValuePercentage, ZeroPercentage, ], False, ), # Numerical ( DataType.bool, [ Count, Mean, Min, Max, Std, ApproximateQuantile, MissingValuePercentage, ZeroPercentage, ], False, ), # Numerical ( DataType.string, [Count, MissingValuePercentage, ApproximateTopK], False, ), # Basic ( DataType.binary, [Count, MissingValuePercentage, ApproximateTopK], False, ), # Basic ( lambda: DataType.temporal("timestamp", unit="us"), [Count, Min, Max, MissingValuePercentage], True, ), # Temporal (pattern matched) ( lambda: DataType.temporal("date32"), [Count, Min, Max, MissingValuePercentage], True, ), # Temporal (pattern matched) ( lambda: DataType.temporal("time64", unit="us"), [Count, Min, Max, MissingValuePercentage], True, ), # Temporal (pattern matched) ], ) def test_default_mappings( self, dtype_factory, expected_agg_types, uses_pattern_matching ): """Test that default mappings return correct aggregators.""" from ray.data.datatype import TypeCategory mapping = _default_dtype_aggregators() dtype = dtype_factory() if uses_pattern_matching: # For pattern-matched types (like temporal), find the matching factory factory = None for mapping_key, mapping_factory in mapping.items(): if isinstance(mapping_key, (TypeCategory, str)) and dtype.is_of( mapping_key ): factory = mapping_factory break assert ( factory is not None ), f"Type {dtype} should match a pattern in the mapping" else: # For exact matches, directly access the mapping assert dtype in mapping factory = mapping[dtype] # Call the factory with a test column to get aggregators aggs = factory("test_col") assert [type(agg) for agg in aggs] == expected_agg_types class TestDatasetSummary: """Test suite for Dataset.summary() method.""" def test_basic_summary(self): """Test basic summary computation.""" ds = ray.data.from_items( [ {"age": 25, "name": "Alice"}, {"age": 30, "name": "Bob"}, ] ) summary = ds.summary() actual = summary.to_pandas() # Verify columns are present assert "age" in actual.columns assert "name" in actual.columns # Check key statistics using rows_same actual_subset = actual[ actual["statistic"].isin(["count", "mean", "min", "max"]) ].copy() actual_subset["age"] = actual_subset["age"].astype(float) actual_subset["name"] = actual_subset["name"].astype(float) expected = pd.DataFrame( { "statistic": ["count", "mean", "min", "max"], "age": [2.0, 27.5, 25.0, 30.0], "name": [2.0, None, None, None], } ) assert rows_same(actual_subset, expected) def test_summary_with_column_filter(self): """Test summary with specific columns.""" ds = ray.data.from_items( [ {"col1": 1, "col2": "a", "col3": 3.5}, ] ) summary = ds.summary(columns=["col1"]) actual = summary.to_pandas() # Check count and mean with rows_same actual_subset = actual[actual["statistic"].isin(["count", "mean"])][ ["statistic", "col1"] ].copy() actual_subset["col1"] = actual_subset["col1"].astype(float) expected = pd.DataFrame( { "statistic": ["count", "mean"], "col1": [1.0, 1.0], } ) assert rows_same(actual_subset, expected) def test_summary_custom_mapping(self): """Test summary with custom dtype aggregation mapping.""" ds = ray.data.from_items([{"value": 10, "other": 20}]) # Only Count and Mean for int64 columns custom_mapping = {DataType.int64(): lambda col: [Count(on=col), Mean(on=col)]} summary = ds.summary(override_dtype_agg_mapping=custom_mapping) actual = summary.to_pandas() # Convert to float for comparison actual["value"] = actual["value"].astype(float) actual["other"] = actual["other"].astype(float) # Columns are sorted alphabetically, so order is: statistic, other, value expected = pd.DataFrame( { "statistic": ["count", "mean"], "other": [1.0, 20.0], "value": [1.0, 10.0], } ) assert rows_same(actual, expected) def test_get_column_stats(self): """Test get_column_stats method.""" ds = ray.data.from_items( [ {"x": 1, "y": 2}, {"x": 3, "y": 4}, ] ) summary = ds.summary() actual = summary.get_column_stats("x") # Verify key statistics with rows_same (checking subset due to mixed types) expected_stats = ["count", "mean", "min", "max"] actual_subset = actual[actual["statistic"].isin(expected_stats)].copy() actual_subset["value"] = actual_subset["value"].astype(float) expected = pd.DataFrame( { "statistic": ["count", "mean", "min", "max"], "value": [2.0, 2.0, 1.0, 3.0], } ) assert rows_same(actual_subset, expected) @pytest.mark.parametrize( "data,column,expected_df", [ ( [{"x": 1}, {"x": 2}, {"x": 3}], "x", pd.DataFrame( { "statistic": ["count", "mean", "min", "max"], "x": [3.0, 2.0, 1.0, 3.0], } ), ), ( [{"y": 10}, {"y": 20}], "y", pd.DataFrame( { "statistic": ["count", "mean", "min", "max"], "y": [2.0, 15.0, 10.0, 20.0], } ), ), ( [{"z": 0}, {"z": 0}, {"z": 1}], "z", pd.DataFrame( { "statistic": ["count", "mean", "min", "max"], "z": [3.0, 1.0 / 3, 0.0, 1.0], } ), ), ], ) def test_summary_statistics_values(self, data, column, expected_df): """Test that computed statistics have correct values.""" ds = ray.data.from_items(data) summary = ds.summary(columns=[column]) actual = summary.to_pandas() # Filter to key statistics and convert to float actual_subset = actual[ actual["statistic"].isin(["count", "mean", "min", "max"]) ][["statistic", column]].copy() actual_subset[column] = actual_subset[column].astype(float) assert rows_same(actual_subset, expected_df) @pytest.mark.parametrize( "data,columns,expected_df", [ # Single numerical column with two values ( [{"x": 10}, {"x": 20}], ["x"], pd.DataFrame( { DatasetSummary.STATISTIC_COLUMN: [ "approx_quantile[0]", "count", "max", "mean", "min", "missing_pct", "std", "zero_pct", ], "x": [20.0, 2.0, 20.0, 15.0, 10.0, 0.0, 5.0, 0.0], } ), ), # Single numerical column with all same values ( [{"y": 5}, {"y": 5}, {"y": 5}], ["y"], pd.DataFrame( { DatasetSummary.STATISTIC_COLUMN: [ "approx_quantile[0]", "count", "max", "mean", "min", "missing_pct", "std", "zero_pct", ], "y": [5.0, 3.0, 5.0, 5.0, 5.0, 0.0, 0.0, 0.0], } ), ), # Multiple numerical columns ( [{"a": 1, "b": 10}, {"a": 3, "b": 30}], ["a", "b"], pd.DataFrame( { DatasetSummary.STATISTIC_COLUMN: [ "approx_quantile[0]", "count", "max", "mean", "min", "missing_pct", "std", "zero_pct", ], "a": [3.0, 2.0, 3.0, 2.0, 1.0, 0.0, 1.0, 0.0], "b": [30.0, 2.0, 30.0, 20.0, 10.0, 0.0, 10.0, 0.0], } ), ), # Column with zeros and missing values ( [{"z": 0}, {"z": 10}, {"z": None}], ["z"], pd.DataFrame( { DatasetSummary.STATISTIC_COLUMN: [ "approx_quantile[0]", "count", "max", "mean", "min", "missing_pct", "std", "zero_pct", ], "z": [10.0, 3.0, 10.0, 5.0, 0.0, 100.0 / 3, 5.0, 50.0], } ), ), # Column with all zeros ( [{"w": 0}, {"w": 0}], ["w"], pd.DataFrame( { DatasetSummary.STATISTIC_COLUMN: [ "approx_quantile[0]", "count", "max", "mean", "min", "missing_pct", "std", "zero_pct", ], "w": [0.0, 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 100.0], } ), ), ], ) def test_summary_full_dataframe(self, data, columns, expected_df): """Test summary with full DataFrame comparison.""" ds = ray.data.from_items(data) summary = ds.summary(columns=columns) actual = summary.to_pandas() # Convert to float for comparison for col in expected_df.columns: if col != DatasetSummary.STATISTIC_COLUMN: expected_df[col] = expected_df[col].astype(float) actual[col] = actual[col].astype(float) assert rows_same(actual, expected_df) def test_summary_multiple_quantiles(self): """Test summary with multiple quantiles.""" ds = ray.data.from_items( [ {"x": 1}, {"x": 2}, {"x": 3}, {"x": 4}, {"x": 5}, ] ) # Create custom mapping with multiple quantiles custom_mapping = { DataType.int64(): lambda col: [ Count(on=col, ignore_nulls=False), Min(on=col, ignore_nulls=True), Max(on=col, ignore_nulls=True), ApproximateQuantile(on=col, quantiles=[0.25, 0.5, 0.75]), ] } summary = ds.summary(columns=["x"], override_dtype_agg_mapping=custom_mapping) actual = summary.to_pandas() # Should have separate rows for each quantile with index-based labels [0], [1], [2] expected = pd.DataFrame( { DatasetSummary.STATISTIC_COLUMN: [ "approx_quantile[0]", "approx_quantile[1]", "approx_quantile[2]", "count", "max", "min", ], "x": [2.0, 3.0, 4.0, 5.0, 5.0, 1.0], } ) # Convert to float for comparison for col in expected.columns: if col != DatasetSummary.STATISTIC_COLUMN: expected[col] = expected[col].astype(float) actual[col] = actual[col].astype(float) assert rows_same(actual, expected) def test_summary_custom_quantiles_and_topk(self): """Test summary with custom ApproximateQuantile and ApproximateTopK values.""" # Create data with numerical and string columns ds = ray.data.from_items( [ {"value": 10, "category": "apple"}, {"value": 20, "category": "banana"}, {"value": 30, "category": "apple"}, {"value": 40, "category": "cherry"}, {"value": 50, "category": "banana"}, {"value": 60, "category": "apple"}, {"value": 70, "category": "date"}, ] ) # Custom mapping with different quantile values and top-k value custom_mapping = { DataType.int64(): lambda col: [ Count(on=col, ignore_nulls=False), ApproximateQuantile(on=col, quantiles=[0.1, 0.5, 0.9]), ], DataType.string(): lambda col: [ Count(on=col, ignore_nulls=False), ApproximateTopK(on=col, k=3), # Top 3 instead of default 10 ], } summary = ds.summary(override_dtype_agg_mapping=custom_mapping) actual = summary.to_pandas() expected_stats = [ "approx_quantile[0]", "approx_quantile[1]", "approx_quantile[2]", "approx_topk[0]", "approx_topk[1]", "approx_topk[2]", "count", ] # Verify all expected statistics are present actual_stats = set(actual[DatasetSummary.STATISTIC_COLUMN].tolist()) assert all(stat in actual_stats for stat in expected_stats) # Helper to get statistic value for a column def get_stat(stat_name, col_name): return actual[actual[DatasetSummary.STATISTIC_COLUMN] == stat_name][ col_name ].iloc[0] # Verify all expected values expected_values = { ("approx_quantile[0]", "value"): 10.0, ("approx_quantile[1]", "value"): 40.0, ("approx_quantile[2]", "value"): 70.0, ("approx_topk[0]", "category"): {"category": "apple", "count": 3}, ("approx_topk[1]", "category"): {"category": "banana", "count": 2}, ("approx_topk[2]", "category"): {"category": "date", "count": 1}, ("count", "value"): 7.0, ("count", "category"): 7, } for (stat, col), expected in expected_values.items(): actual_value = get_stat(stat, col) assert ( actual_value == expected ), f"{stat}[{col}] = {actual_value} != {expected}" if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))