fix: Fix default AutoML Forecasting transformations list.

PiperOrigin-RevId: 526734524

google/cloud/aiplatform/training_jobs.py

@@ -2438,7 +2438,9 @@ def _run(
  (
  self._column_transformations,
  column_names,
- ) = dataset._get_default_column_transformations(target_column)
+ ) = column_transformations_utils.get_default_column_transformations(
+ dataset=dataset, target_column=target_column
+ )
 
  _LOGGER.info(
  "The column transformation of type 'auto' was set for the following columns: %s."

google/cloud/aiplatform/utils/column_transformations_utils.py

@@ -16,7 +16,6 @@
 #
 
 from typing import Dict, List, Optional, Tuple
-import warnings
 
 from google.cloud.aiplatform import datasets
 
@@ -51,9 +50,9 @@ def get_default_column_transformations(
 
 
 def validate_and_get_column_transformations(
- column_specs: Optional[Dict[str, str]],
- column_transformations: Optional[List[Dict[str, Dict[str, str]]]],
-) -> List[Dict[str, Dict[str, str]]]:
+ column_specs: Optional[Dict[str, str]] = None,
+ column_transformations: Optional[List[Dict[str, Dict[str, str]]]] = None,
+) -> Optional[List[Dict[str, Dict[str, str]]]]:
  """Validates column specs and transformations, then returns processed transformations.
 
  Args:
@@ -91,21 +90,13 @@ def validate_and_get_column_transformations(
  # user populated transformations
  if column_transformations is not None and column_specs is not None:
  raise ValueError(
- "Both column_transformations and column_specs were passed. Only one is allowed."
+ "Both column_transformations and column_specs were passed. Only "
+ "one is allowed."
  )
- if column_transformations is not None:
- warnings.simplefilter("always", DeprecationWarning)
- warnings.warn(
- "consider using column_specs instead. column_transformations will be deprecated in the future.",
- DeprecationWarning,
- stacklevel=2,
- )
-
- return column_transformations
  elif column_specs is not None:
  return [
  {transformation: {"column_name": column_name}}
  for column_name, transformation in column_specs.items()
  ]
  else:
- return None
+ return column_transformations

tests/unit/aiplatform/test_automl_forecasting_training_jobs.py

@@ -1294,3 +1294,41 @@ def test_run_call_pipeline_if_set_additional_experiments_probabilistic_inference
  training_pipeline=true_training_pipeline,
  timeout=None,
  )
+
+ def test_automl_forecasting_with_no_transformations(
+ self,
+ mock_pipeline_service_create,
+ mock_pipeline_service_get,
+ mock_dataset_time_series,
+ mock_model_service_get,
+ ):
+ aiplatform.init(project=_TEST_PROJECT)
+ job = training_jobs.AutoMLForecastingTrainingJob(
+ display_name=_TEST_DISPLAY_NAME,
+ optimization_objective=_TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME,
+ )
+ mock_dataset_time_series.column_names = [
+ "a",
+ "b",
+ _TEST_TRAINING_TARGET_COLUMN,
+ ]
+ job.run(
+ dataset=mock_dataset_time_series,
+ predefined_split_column_name=_TEST_PREDEFINED_SPLIT_COLUMN_NAME,
+ target_column=_TEST_TRAINING_TARGET_COLUMN,
+ time_column=_TEST_TRAINING_TIME_COLUMN,
+ time_series_identifier_column=_TEST_TRAINING_TIME_SERIES_IDENTIFIER_COLUMN,
+ unavailable_at_forecast_columns=_TEST_TRAINING_UNAVAILABLE_AT_FORECAST_COLUMNS,
+ available_at_forecast_columns=_TEST_TRAINING_AVAILABLE_AT_FORECAST_COLUMNS,
+ forecast_horizon=_TEST_TRAINING_FORECAST_HORIZON,
+ data_granularity_unit=_TEST_TRAINING_DATA_GRANULARITY_UNIT,
+ data_granularity_count=_TEST_TRAINING_DATA_GRANULARITY_COUNT,
+ model_display_name=_TEST_MODEL_DISPLAY_NAME,
+ time_series_attribute_columns=_TEST_TRAINING_TIME_SERIES_ATTRIBUTE_COLUMNS,
+ context_window=_TEST_TRAINING_CONTEXT_WINDOW,
+ budget_milli_node_hours=_TEST_TRAINING_BUDGET_MILLI_NODE_HOURS,
+ )
+ assert job._column_transformations == [
+ {"auto": {"column_name": "a"}},
+ {"auto": {"column_name": "b"}},
+ ]

tests/unit/aiplatform/test_utils.py

@@ -35,7 +35,9 @@
 from google.cloud import storage
 from google.cloud.aiplatform import compat, utils
 from google.cloud.aiplatform.compat.types import pipeline_failure_policy
+from google.cloud.aiplatform import datasets
 from google.cloud.aiplatform.utils import (
+ column_transformations_utils,
  gcs_utils,
  pipeline_utils,
  prediction_utils,
@@ -485,6 +487,51 @@ def test_timestamped_unique_name():
  assert re.match(r"\d{4}-\d{2}-\d{2}-\d{2}-\d{2}-\d{2}-.{5}", name)
 
 
+class TestColumnTransformationsUtils:
+
+ column_transformations = [
+ {"auto": {"column_name": "a"}},
+ {"auto": {"column_name": "b"}},
+ ]
+ column_specs = {"a": "auto", "b": "auto"}
+
+ def test_get_default_column_transformations(self):
+ ds = mock.MagicMock(datasets.TimeSeriesDataset)
+ ds.column_names = ["a", "b", "target"]
+ (
+ transforms,
+ columns,
+ ) = column_transformations_utils.get_default_column_transformations(
+ dataset=ds, target_column="target"
+ )
+ assert transforms == [
+ {"auto": {"column_name": "a"}},
+ {"auto": {"column_name": "b"}},
+ ]
+ assert columns == ["a", "b"]
+
+ def test_validate_transformations_with_multiple_configs(self):
+ with pytest.raises(ValueError):
+ (
+ column_transformations_utils.validate_and_get_column_transformations(
+ column_transformations=self.column_transformations,
+ column_specs=self.column_specs,
+ )
+ )
+
+ def test_validate_transformations_with_column_specs(self):
+ actual = column_transformations_utils.validate_and_get_column_transformations(
+ column_specs=self.column_specs
+ )
+ assert actual == self.column_transformations
+
+ def test_validate_transformations_with_column_transformations(self):
+ actual = column_transformations_utils.validate_and_get_column_transformations(
+ column_transformations=self.column_transformations
+ )
+ assert actual == self.column_transformations
+
+
 @pytest.mark.usefixtures("google_auth_mock")
 class TestGcsUtils:
  def test_upload_to_gcs(self, json_file, mock_storage_blob_upload_from_filename):