Allow me to turn off or control any subsampling done within the quality report (#792)

Author: R-Palazzo

sdmetrics/reports/base_report.py

@@ -14,6 +14,7 @@
 import tqdm
 
 from sdmetrics._utils_metadata import _convert_datetime_column, _validate_metadata
+from sdmetrics.reports.utils import DEFAULT_NUM_ROWS_SUBSAMPLE
 from sdmetrics.visualization import set_plotly_config
 
 
@@ -27,6 +28,7 @@ def __init__(self):
         self._overall_score = None
         self.is_generated = False
         self._properties = {}
+        self.num_rows_subsample = DEFAULT_NUM_ROWS_SUBSAMPLE
         self.report_info = {
             'report_type': self.__class__.__name__,
             'generated_date': None,
@@ -163,6 +165,7 @@ def generate(self, real_data, synthetic_data, metadata, verbose=True):
                     f'({ind + 1}/{len(self._properties)}) Evaluating {property_name}'
                 )
 
+            self._properties[property_name].num_rows_subsample = self.num_rows_subsample
             score = self._properties[property_name].get_score(
                 real_data, synthetic_data, metadata, progress_bar=progress_bar
             )

sdmetrics/reports/multi_table/_properties/base.py

@@ -3,6 +3,8 @@
 import numpy as np
 import pandas as pd
 
+from sdmetrics.reports.utils import DEFAULT_NUM_ROWS_SUBSAMPLE
+
 
 class BaseMultiTableProperty:
     """Base class for multi table properties.
@@ -26,6 +28,7 @@ def __init__(self):
         self._properties = {}
         self.is_computed = False
         self.details = pd.DataFrame()
+        self.num_rows_subsample = DEFAULT_NUM_ROWS_SUBSAMPLE
 
     def _get_num_iterations(self, metadata):
         """Get the number of iterations for the property."""

sdmetrics/reports/single_table/_properties/base.py

@@ -2,6 +2,8 @@
 
 import pandas as pd
 
+from sdmetrics.reports.utils import DEFAULT_NUM_ROWS_SUBSAMPLE
+
 
 class BaseSingleTableProperty:
     """Base class for single table properties.
@@ -14,6 +16,7 @@ class BaseSingleTableProperty:
 
     def __init__(self):
         self.details = pd.DataFrame()
+        self.num_rows_subsample = DEFAULT_NUM_ROWS_SUBSAMPLE
 
     def _compute_average(self):
         """Average the scores for each column."""

sdmetrics/reports/single_table/_properties/column_pair_trends.py

@@ -10,8 +10,6 @@
 from sdmetrics.reports.single_table._properties import BaseSingleTableProperty
 from sdmetrics.reports.utils import PlotConfig
 
-DEFAULT_NUM_ROWS_SUBSAMPLE = 50000
-
 
 class ColumnPairTrends(BaseSingleTableProperty):
     """Column pair trends property.
@@ -30,6 +28,7 @@ class ColumnPairTrends(BaseSingleTableProperty):
     }
 
     def __init__(self):
+        super().__init__()
         self._columns_datetime_conversion_failed = {}
         self._columns_discretization_failed = {}
 
@@ -276,10 +275,12 @@ def _generate_details(
             )
 
             metric_params = {}
-            if (metric == ContingencySimilarity) and (
-                max(len(col_real), len(col_synthetic)) > DEFAULT_NUM_ROWS_SUBSAMPLE
+            if (
+                self.num_rows_subsample
+                and (metric == ContingencySimilarity)
+                and (max(len(col_real), len(col_synthetic)) > self.num_rows_subsample)
             ):
-                metric_params['num_rows_subsample'] = DEFAULT_NUM_ROWS_SUBSAMPLE
+                metric_params['num_rows_subsample'] = self.num_rows_subsample
 
             try:
                 error = self._preprocessing_failed(

sdmetrics/reports/utils.py

@@ -17,6 +17,7 @@
 
 CONTINUOUS_SDTYPES = ['numerical', 'datetime']
 DISCRETE_SDTYPES = ['categorical', 'boolean']
+DEFAULT_NUM_ROWS_SUBSAMPLE = 50000
 
 
 class PlotConfig:

tests/integration/reports/single_table/test_quality_report.py

@@ -77,6 +77,7 @@ def test_report_end_to_end(self):
             key: val for key, val in metadata['columns'].items() if key in column_names
         }
         report = QualityReport()
+        report.num_rows_subsample = None
 
         # Run
         generate_start_time = time.time()
@@ -141,7 +142,8 @@ def test_report_end_to_end(self):
             report.get_details('Column Pair Trends'), expected_details_cpt
         )
         assert report.get_score() == 0.8393750143888287
-
+        assert report._properties['Column Shapes'].num_rows_subsample is None
+        assert report._properties['Column Pair Trends'].num_rows_subsample is None
         report_info = report.get_info()
         assert report_info == report.report_info
 
@@ -183,6 +185,8 @@ def test_with_large_dataset(self):
         # Assert
         cpt_report_1 = report_1.get_properties().iloc[1]['Score']
         cpt_report_2 = report_2.get_properties().iloc[1]['Score']
+        assert report_1._properties['Column Pair Trends'].num_rows_subsample == 50000
+        assert report_2._properties['Column Pair Trends'].num_rows_subsample == 50000
         assert score_1_run_1 != score_1_run_2
         assert np.isclose(score_1_run_1, score_1_run_2, atol=0.001)
         assert np.isclose(report_2.get_score(), score_1_run_1, atol=0.001)

tests/unit/reports/multi_table/test_base_multi_table_report.py

@@ -8,6 +8,7 @@
 
 from sdmetrics.demos import load_demo
 from sdmetrics.reports.multi_table.base_multi_table_report import BaseMultiTableReport
+from sdmetrics.reports.utils import DEFAULT_NUM_ROWS_SUBSAMPLE
 
 
 class TestBaseReport:
@@ -21,6 +22,7 @@ def test__init__(self):
         assert report.is_generated is False
         assert report._properties == {}
         assert report.table_names == []
+        assert report.num_rows_subsample == DEFAULT_NUM_ROWS_SUBSAMPLE
 
     def test__validate_data_format(self):
         """Test the ``_validate_data_format`` method.

tests/unit/reports/test_base_report.py

@@ -9,9 +9,21 @@
 
 from sdmetrics.demos import load_demo
 from sdmetrics.reports.base_report import BaseReport
+from sdmetrics.reports.utils import DEFAULT_NUM_ROWS_SUBSAMPLE
 
 
 class TestBaseReport:
+    def test__init__(self):
+        """Test the initialization of the BaseReport class."""
+        # Run
+        base_report = BaseReport()
+
+        # Assert
+        assert base_report._overall_score is None
+        assert not base_report.is_generated
+        assert base_report._properties == {}
+        assert base_report.num_rows_subsample == DEFAULT_NUM_ROWS_SUBSAMPLE
+
     def test__validate_data_format(self):
         """Test the ``_validate_data_format`` method.
 
@@ -268,6 +280,7 @@ def test_generate(self, version_mock, time_mock, datetime_mock):
         version_mock.return_value = 'version'
 
         base_report = BaseReport()
+        base_report.num_rows_subsample = 1000
         mock_validate = Mock()
         mock__print_results = Mock()
         base_report._print_results = mock__print_results
@@ -292,9 +305,11 @@ def test_generate(self, version_mock, time_mock, datetime_mock):
         base_report._properties['Property 1'].get_score.assert_called_with(
             real_data, synthetic_data, metadata, progress_bar=None
         )
+        assert base_report._properties['Property 1'].num_rows_subsample == 1000
         base_report._properties['Property 2'].get_score.assert_called_with(
             real_data, synthetic_data, metadata, progress_bar=None
         )
+        assert base_report._properties['Property 2'].num_rows_subsample == 1000
         expected_info = {
             'report_type': 'BaseReport',
             'generated_date': '2020-01-05',