Avoid using multiprocessing by default when generating statistics over a

dataframe.

PiperOrigin-RevId: 236177889

Author: paulgc17

RELEASE.md

@@ -9,11 +9,12 @@
 
 ## Bug Fixes and Other Changes
 * Expand unit test coverage.
-
 * Modify validation logic to raise `SCHEMA_MISSING_COLUMN` anomaly when
   observing a feature with no stats.
 * Add utility functions `write_stats_text` and `load_stats_text` to write and
   load DatasetFeatureStatisticsList protos.
+* Avoid using multiprocessing by default when generating statistics over a
+  dataframe.
 * Depends on `joblib>=0.12,<1`.
 * Requires pre-installed `tensorflow>=1.13,<2`.
 

tensorflow_data_validation/utils/stats_gen_lib.py

@@ -173,7 +173,7 @@ def generate_statistics_from_csv(
 def generate_statistics_from_dataframe(
     dataframe,
     stats_options = options.StatsOptions(),
-    n_jobs = multiprocessing.cpu_count()
+    n_jobs = 1
 ):
   """Compute data statistics for the input pandas DataFrame.
 
@@ -183,7 +183,8 @@ def generate_statistics_from_dataframe(
   Args:
     dataframe: Input pandas DataFrame.
     stats_options: Options for generating data statistics.
-    n_jobs: Number of processes to run.
+    n_jobs: Number of processes to run (defaults to 1). If -1 is provided,
+      uses the same number of processes as the number of CPU cores.
 
   Returns:
     A DatasetFeatureStatisticsList proto.
@@ -193,7 +194,14 @@ def generate_statistics_from_dataframe(
                     'pandas DataFrame.'.format(type(dataframe).__name__))
 
   stats_generators = stats_impl.get_generators(stats_options, in_memory=True)
+  if n_jobs < -1 or n_jobs == 0:
+    raise ValueError('Invalid n_jobs parameter {}. Should be either '
+                     ' -1 or >= 1.'.format(n_jobs))
+
+  if n_jobs == -1:
+    n_jobs = multiprocessing.cpu_count()
   n_jobs = max(min(n_jobs, multiprocessing.cpu_count()), 1)
+
   if n_jobs == 1:
     merged_partial_stats = _generate_partial_statistics_from_df(
         dataframe, stats_options, stats_generators)

tensorflow_data_validation/utils/stats_gen_lib_test.py

@@ -565,6 +565,28 @@ def test_stats_gen_with_dataframe(self):
     test_util.assert_dataset_feature_stats_proto_equal(
         self, result.datasets[0], expected_result.datasets[0])
 
+  def test_stats_gen_with_dataframe_invalid_njobs_zero(self):
+    records, _, _ = self._get_csv_test(delimiter=',', with_header=True)
+    input_data_path = self._write_records_to_csv(records, self._get_temp_dir(),
+                                                 'input_data.csv')
+    dataframe = pd.read_csv(input_data_path)
+    with self.assertRaisesRegexp(
+        ValueError, 'Invalid n_jobs parameter.*'):
+      _ = stats_gen_lib.generate_statistics_from_dataframe(
+          dataframe=dataframe,
+          stats_options=self._default_stats_options, n_jobs=0)
+
+  def test_stats_gen_with_dataframe_invalid_njobs_negative(self):
+    records, _, _ = self._get_csv_test(delimiter=',', with_header=True)
+    input_data_path = self._write_records_to_csv(records, self._get_temp_dir(),
+                                                 'input_data.csv')
+    dataframe = pd.read_csv(input_data_path)
+    with self.assertRaisesRegexp(
+        ValueError, 'Invalid n_jobs parameter.*'):
+      _ = stats_gen_lib.generate_statistics_from_dataframe(
+          dataframe=dataframe,
+          stats_options=self._default_stats_options, n_jobs=-2)
+
 
 if __name__ == '__main__':
   absltest.main()