Update data loaders. Now label column filter can be an integer or a l…

…ist of integers.

PiperOrigin-RevId: 635061726

CHANGELOG.md

@@ -23,14 +23,19 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
-## [0.2.0] - 2022-05-05
+## [0.2.1] - 2024-05-18
+
+* Updates to data loaders. Label column filter can now be a list of integers.
+
+## [0.2.0] - 2024-05-05
 
 * Add PyPi support. Minor reorganization of repository.
 
 ## [0.1.0] - 2024-04-17
 
 * Initial release
 
-[Unreleased]: https://github.com/google-research/spade_anomaly_detection/compare/v0.2.0...HEAD
+[Unreleased]: https://github.com/google-research/spade_anomaly_detection/compare/v0.2.1...HEAD
+[0.2.1]: https://github.com/google-research/spade_anomaly_detection/compare/v0.2.0...v0.2.1
 [0.2.0]: https://github.com/google-research/spade_anomaly_detection/compare/v0.1.0...v0.2.0
 [0.1.0]: https://github.com/google-research/spade_anomaly_detection/releases/tag/v0.1.0

pyproject.toml

@@ -2,7 +2,6 @@
 # Project metadata. Available keys are documented at:
 # https://packaging.python.org/en/latest/specifications/declaring-project-metadata
 name = "spade_anomaly_detection"
-version = "0.2.0"
 description = "Semi-supervised Pseudo Labeler Anomaly Detection with Ensembling (SPADE) is a semi-supervised anomaly detection method that uses an ensemble of one class classifiers as the pseudo-labelers and supervised classifiers to achieve state of the art results especially on datasets with distribution mismatch between labeled and unlabeled samples."
 readme = "README.md"
 requires-python = ">=3.8"
@@ -35,7 +34,7 @@ dependencies = [
 ]
 
 # `version` is automatically set by flit to use `spade_anomaly_detection.__version__`
-# dynamic = ["version"]
+dynamic = ["version"]
 
 [project.urls]
 homepage = "https://github.com/google-research/spade_anomaly_detection"

spade_anomaly_detection/__init__.py

@@ -31,4 +31,4 @@
 
 # A new PyPI release will be pushed every time `__version__` is increased.
 # When changing this, also update the CHANGELOG.md.
-__version__ = '0.2.0'
+__version__ = '0.2.1'

spade_anomaly_detection/data_loader.py

@@ -328,8 +328,9 @@ def load_tf_dataset_from_bigquery(
  where_statements: Optional[List[str]] = None,
  ignore_columns: Optional[Sequence[str]] = None,
  batch_size: Optional[int] = None,
- label_column_filter_value: Optional[int] = None,
+ label_column_filter_value: Optional[int | list[int]] = None,
  convert_features_to_float64: bool = False,
+ page_size: Optional[int] = None,
  ) -> tf.data.Dataset:
  """Loads a TensorFlow dataset from a BigQuery Table.
 
@@ -346,10 +347,13 @@ def load_tf_dataset_from_bigquery(
  dataset is not batched. In this case, when iterating through the
  dataset, it will yield one record per call instead of a batch of
  records.
- label_column_filter_value: An integer used when filtering the label column
- values. No value will result in all data returned from the table.
+ label_column_filter_value: An integer or list of integers used when
+ filtering the label column values. No value will result in all data
+ returned from the table.
  convert_features_to_float64: Set to True to cast the contents of the
  features columns to float64.
+ page_size: the pagination size to use when retrieving data from BigQuery.
+ A large value can result in fewer BQ calls, hence time savings.
 
  Returns:
  A TensorFlow dataset.
@@ -362,7 +366,15 @@ def load_tf_dataset_from_bigquery(
  where_statements = (
  list() if where_statements is None else where_statements
  )
- where_statements.append(f'{label_col_name} = {label_column_filter_value}')
+ if isinstance(label_column_filter_value, int):
+ where_statements.append(
+ f'{label_col_name} = {label_column_filter_value}'
+ )
+ else:
+ where_statements.append(
+ f'CAST({label_col_name} AS INT64) IN '
+ f'UNNEST({label_column_filter_value})'
+ )
 
  if ignore_columns is not None:
  metadata_builder = feature_metadata.BigQueryMetadataBuilder(
@@ -373,8 +385,10 @@ def load_tf_dataset_from_bigquery(
  )
 
  if where_statements:
- metadata_retrieval_options = feature_metadata.MetadataRetrievalOptions(
- where_clauses=where_statements
+ metadata_retrieval_options = (
+ feature_metadata.MetadataRetrievalOptions.get_none(
+ where_clauses=where_statements
+ )
  )
 
  tf_dataset, metadata = bq_dataset.get_dataset_and_metadata_for_table(
@@ -384,7 +398,12 @@ def load_tf_dataset_from_bigquery(
  drop_remainder=True,
  metadata_options=metadata_retrieval_options,
  metadata_builder=metadata_builder,
+ page_size=page_size,
  )
+ options = tf.data.Options()
+ # Avoid a large warning output by TF Dataset.
+ options.deterministic = False
+ tf_dataset = tf_dataset.with_options(options)
 
  self.input_feature_metadata = metadata
 

spade_anomaly_detection/data_loader_test.py

@@ -341,7 +341,9 @@ def test_load_bigquery_dataset_unlabeled_value(self, metadata_mock):
  label_column_filter_value=label_column_filter_value,
  )
 
- mock_metadata_call_actual = metadata_mock.call_args.kwargs['where_clauses']
+ mock_metadata_call_actual = metadata_mock.get_none.call_args.kwargs[
+ 'where_clauses'
+ ]
 
  # Ensure that a where statement was not created when we don't pass in label
  # values.
@@ -399,7 +401,7 @@ def test_where_statement_construction_no_error(self, mock_metadata):
 
  self.assertListEqual(
  expected_where_statements,
- mock_metadata.call_args.kwargs['where_clauses'],
+ mock_metadata.get_none.call_args.kwargs['where_clauses'],
  )
 
  @mock.patch.object(
@@ -419,7 +421,7 @@ def test_where_statements_with_label_filter_no_error(self, mock_metadata):
 
  self.assertListEqual(
  where_statements,
- mock_metadata.call_args.kwargs['where_clauses'],
+ mock_metadata.get_none.call_args.kwargs['where_clauses'],
  )
 
  @mock.patch.object(feature_metadata, 'BigQueryMetadataBuilder', autospec=True)

spade_anomaly_detection/data_utils/bq_dataset.py

@@ -721,6 +721,7 @@ def get_dataset_and_metadata_for_table(
  batch_size: int = 64,
  with_mask: bool = _WITH_MASK_DEFAULT,
  drop_remainder: bool = False,
+ page_size: Optional[int] = None,
 ) -> Tuple[tf.data.Dataset, feature_metadata.BigQueryTableMetadata]:
  """Gets the metadata and dataset for a BigQuery table.
 
@@ -735,6 +736,8 @@ def get_dataset_and_metadata_for_table(
  with_mask: Whether the dataset should be returned with a mask format. For
  more information see get_bigquery_dataset.
  drop_remainder: If true no partial batches will be yielded.
+ page_size: the pagination size to use when retrieving data from BigQuery. A
+ large value can result in fewer BQ calls, hence time savings.
 
  Returns:
  A tuple of the output dataset and metadata for the specified table.
@@ -783,6 +786,7 @@ def get_dataset_and_metadata_for_table(
  cache_location=None,
  where_clauses=metadata_options.where_clauses,
  drop_remainder=drop_remainder,
+ page_size=page_size,
  )
 
  return dataset, all_metadata

spade_anomaly_detection/data_utils/bq_dataset_test.py

@@ -683,6 +683,7 @@ def _create_rand_df():
  self.assertIsInstance(output_dataset, tf.data.Dataset)
 
  # Loop through the batches and make sure they all match.
+ epoch = None
  for epoch in range(1, 3):
  batch_index = 0
  for cur_batch in output_dataset:
@@ -760,7 +761,7 @@ def test_get_dataset_and_metadata_for_table_path(
  mock_bq_storage_client = mock.create_autospec(
  bigquery_storage.BigQueryReadClient, spec_set=True, instance=True)
  metadata_options = feature_metadata.MetadataRetrievalOptions.get_all()
- batch_sie = 128
+ batch_size = 128
  with_mask = False
 
  output_dataset, output_metadata = (
@@ -769,7 +770,7 @@ def test_get_dataset_and_metadata_for_table_path(
  bigquery_client=mock_bq_client,
  bigquery_storage_client=mock_bq_storage_client,
  metadata_options=metadata_options,
- batch_size=batch_sie,
+ batch_size=batch_size,
  with_mask=with_mask,
  )
  )
@@ -780,11 +781,12 @@ def test_get_dataset_and_metadata_for_table_path(
  get_metadata_mock.return_value,
  mock_bq_client,
  bqstorage_client=mock_bq_storage_client,
- batch_size=batch_sie,
+ batch_size=batch_size,
  with_mask=with_mask,
  cache_location=None,
  where_clauses=(),
  drop_remainder=False,
+ page_size=None,
  )
 
  self.assertEqual(output_dataset, get_bigquery_dataset_mock.return_value)
@@ -832,6 +834,7 @@ def test_get_dataset_and_metadata_for_table_parts_defaults(
  cache_location=None,
  where_clauses=(),
  drop_remainder=False,
+ page_size=None,
  )
 
  self.assertEqual(output_dataset, get_bigquery_dataset_mock.return_value)