seed dataset statistics limited to general stats (#32)

Author: johnnygreco

src/data_designer/config/analysis/column_statistics.py

@@ -148,13 +148,8 @@ def create_report_row_data(self) -> dict[str, str]:
 
 
 class SeedDatasetColumnStatistics(GeneralColumnStatistics):
-    distribution_type: ColumnDistributionType
-    distribution: Optional[Union[CategoricalDistribution, NumericalDistribution, MissingValue]]
     column_type: Literal[DataDesignerColumnType.SEED_DATASET.value] = DataDesignerColumnType.SEED_DATASET.value
 
-    def create_report_row_data(self) -> dict[str, str]:
-        return self._general_display_row
-
 
 class ExpressionColumnStatistics(GeneralColumnStatistics):
     column_type: Literal[DataDesignerColumnType.EXPRESSION.value] = DataDesignerColumnType.EXPRESSION.value

src/data_designer/engine/analysis/column_statistics.py

@@ -24,7 +24,6 @@
     calculate_general_column_info,
     calculate_token_stats,
     calculate_validation_column_info,
-    determine_column_distribution_type,
 )
 
 logger = logging.getLogger(__name__)
@@ -105,18 +104,7 @@ def calculate_sampler_distribution(self) -> dict[str, Any]:
         )
 
 
-class SeedDatasetColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
-    def calculate_seed_dataset_distribution(self) -> dict[str, Any]:
-        dist_type = determine_column_distribution_type(self.df[self.column_config.name])
-        make_dist = dist_type in [ColumnDistributionType.CATEGORICAL, ColumnDistributionType.NUMERICAL]
-        return (
-            calculate_column_distribution(self.column_config, self.df, dist_type)
-            if make_dist
-            else {
-                "distribution_type": dist_type,
-                "distribution": None,
-            }
-        )
+class SeedDatasetColumnStatisticsCalculator(GeneralColumnStatisticsCalculator): ...
 
 
 class ValidationColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):

src/data_designer/engine/analysis/utils/column_statistics_calculations.py

@@ -9,8 +9,6 @@
 
 import numpy as np
 import pandas as pd
-from pandas import Series
-from pandas.core.dtypes.common import is_integer_dtype, is_numeric_dtype
 import pyarrow as pa
 import tiktoken
 
@@ -180,67 +178,6 @@ def convert_pyarrow_dtype_to_simple_dtype(pyarrow_dtype: pa.DataType) -> str:
     return pyarrow_dtype_str
 
 
-def determine_column_distribution_type(column: Series) -> ColumnDistributionType:
-    """Based on the logic used by Gretel's SQS report to determine column data type."""
-    if len(column) == 0:
-        return ColumnDistributionType.OTHER
-
-    if isinstance(column.iloc[0], np.ndarray):
-        return ColumnDistributionType.OTHER
-
-    if isinstance(column.iloc[0], dict):
-        return ColumnDistributionType.OTHER
-
-    try:
-        non_na_data = column.dropna()
-        non_na_count = int(non_na_data.count())
-        unique_count = int(non_na_data.nunique())
-    except Exception:
-        column = column.astype(str)
-        non_na_data = column.dropna()
-        non_na_count = int(non_na_data.count())
-        unique_count = int(non_na_data.nunique())
-
-    if non_na_count == 0:
-        return ColumnDistributionType.OTHER
-
-    if is_numeric_dtype(non_na_data.dtype):
-        # Float values that are within 1e-8 of an integer are considered integers
-        # Floats are considered numerical.
-        if not np.allclose(non_na_data, non_na_data.astype(int), atol=1e-8):
-            return ColumnDistributionType.NUMERICAL
-        # We can visualize numeric data with histograms, but we will not use it for diversity calculations
-        min_value = int(non_na_data.min())
-        if unique_count <= 10 and min_value >= 0:
-            return ColumnDistributionType.CATEGORICAL
-        if unique_count == non_na_count and is_integer_dtype(non_na_data.dtype):
-            # All unique integer values, potentially an ID column
-            return ColumnDistributionType.OTHER
-        return ColumnDistributionType.NUMERICAL
-
-    # Check if the column is a date-like column before checking for categorical or text columns.
-    try:
-        pd.to_datetime(non_na_data, format="%Y-%m-%d")
-        return ColumnDistributionType.OTHER
-    except Exception:
-        pass
-
-    diff = non_na_count - unique_count
-    diff_percent = diff / non_na_count
-    if diff_percent >= 0.9 or (diff_percent >= 0.7 and len(non_na_data) <= 50):
-        return ColumnDistributionType.CATEGORICAL
-
-    space_count = sum(str(entry).strip().count(" ") for entry in non_na_data)
-    if space_count / non_na_count > TEXT_FIELD_AVG_SPACE_COUNT_THRESHOLD:
-        return ColumnDistributionType.TEXT
-
-    if pd.api.types.is_string_dtype(non_na_data.dtype) and unique_count <= 10:
-        # Check for string columns with a small number of unique values (categorical)
-        return ColumnDistributionType.CATEGORICAL
-
-    return ColumnDistributionType.OTHER
-
-
 def ensure_hashable(x: Any) -> str:
     """
     Makes a best effort turn known unhashable types to a hashable

tests/config/analysis/test_column_statistics.py

@@ -19,7 +19,6 @@
     NumericalDistribution,
     SamplerColumnStatistics,
     SamplerType,
-    SeedDatasetColumnStatistics,
     ValidationColumnStatistics,
 )
 
@@ -190,22 +189,6 @@ def test_sampler_column_statistics(stub_general_stats_args_with_valid_values, st
     }
 
 
-def test_seed_dataset_column_statistics(stub_general_stats_args_with_valid_values, stub_categorical_distribution):
-    seed_dataset_column_statistics = SeedDatasetColumnStatistics(
-        **stub_general_stats_args_with_valid_values,
-        distribution_type=ColumnDistributionType.CATEGORICAL,
-        distribution=stub_categorical_distribution,
-    )
-    assert seed_dataset_column_statistics.column_type == "seed-dataset"
-    assert seed_dataset_column_statistics.distribution_type == ColumnDistributionType.CATEGORICAL
-    assert isinstance(seed_dataset_column_statistics.distribution, CategoricalDistribution)
-    assert seed_dataset_column_statistics.create_report_row_data() == {
-        "column name": "test",
-        "number unique values": "10 (10.0%)",
-        "data type": "str",
-    }
-
-
 def test_validation_column_statistics_with_missing_values(stub_general_stats_args_with_missing_values):
     validation_column_statistics = ValidationColumnStatistics(
         **stub_general_stats_args_with_missing_values,

tests/engine/analysis/test_column_statistics_calculator.py

@@ -71,25 +71,3 @@ def test_sampler_column_statistics(stub_df, column_configs):
                 assert isinstance(stats.distribution.mean, float)
                 assert isinstance(stats.distribution.stddev, float)
                 assert isinstance(stats.distribution.median, float)
-
-
-def test_seed_dataset_column_statistics(stub_df, column_configs):
-    for column_config in column_configs:
-        if column_config.column_type == DataDesignerColumnType.SEED_DATASET:
-            column_config_with_df = ColumnConfigWithDataFrame(column_config=column_config, df=stub_df)
-            stats = get_column_statistics_calculator(column_config.column_type)(
-                column_config_with_df=column_config_with_df
-            ).calculate()
-            assert stats.column_name == column_config.name
-            assert stats.column_type == column_config.column_type
-            if stats.distribution_type == ColumnDistributionType.CATEGORICAL:
-                assert hasattr(stats.distribution, "histogram")
-                assert isinstance(stats.distribution.most_common_value, (int, str))
-                assert isinstance(stats.distribution.least_common_value, (int, str))
-            elif stats.distribution_type == ColumnDistributionType.NUMERICAL:
-                assert not hasattr(stats.distribution, "histogram")
-                assert isinstance(stats.distribution.min, (int, float))
-                assert isinstance(stats.distribution.max, (int, float))
-                assert isinstance(stats.distribution.mean, float)
-                assert isinstance(stats.distribution.stddev, float)
-                assert isinstance(stats.distribution.median, float)

tests/engine/analysis/utils/test_column_statistics_calculations.py

@@ -25,7 +25,6 @@
     calculate_prompt_token_stats,
     calculate_validation_column_info,
     convert_pyarrow_dtype_to_simple_dtype,
-    determine_column_distribution_type,
     ensure_boolean,
     ensure_hashable,
 )
@@ -244,29 +243,6 @@ def test_convert_pyarrow_dtype_to_simple_dtype():
     assert convert_pyarrow_dtype_to_simple_dtype(unknown_type) == str(unknown_type)
 
 
-def test_determine_column_distribution_type():
-    assert determine_column_distribution_type(pd.Series([])) == ColumnDistributionType.OTHER
-    assert determine_column_distribution_type(pd.Series([{"a": 1}, {"b": 2}])) == ColumnDistributionType.OTHER
-    assert (
-        determine_column_distribution_type(pd.Series([np.array([1, 2, 3]), np.array([4, 5, 6])]))
-        == ColumnDistributionType.OTHER
-    )
-    assert determine_column_distribution_type(pd.Series([1, 2, 1, 3, 1, 2])) == ColumnDistributionType.CATEGORICAL
-    assert determine_column_distribution_type(pd.Series([1.1, 2.2, 3.3, 4.4, 5.5])) == ColumnDistributionType.NUMERICAL
-    assert (
-        determine_column_distribution_type(pd.Series(["A", "A", "C", "C", "A", "B"]))
-        == ColumnDistributionType.CATEGORICAL
-    )
-    assert (
-        determine_column_distribution_type(pd.Series(["This is a long text", "Another long text with spaces"]))
-        == ColumnDistributionType.TEXT
-    )
-    assert (
-        determine_column_distribution_type(pd.Series(["2023-01-01", "2023-01-02", "2023-01-03"]))
-        == ColumnDistributionType.OTHER
-    )
-
-
 def test_prepare_number_for_reporting():
     assert prepare_number_for_reporting(5, int) == 5
     assert isinstance(prepare_number_for_reporting(5, int), int)