Merge pull request #3808 from AI-Hypercomputer:aireen/grain_error

PiperOrigin-RevId: 910999131

Author: Google-ML-Automation

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -556,14 +556,20 @@ def map(self, element):
     example.ParseFromString(element)
     features = example.features.feature
 
+    missing = [c for c in self.data_columns if c not in features]
+    if missing:
+      raise ValueError(
+          f"Column {missing} not found in dataset. Available columns: {sorted(features.keys())}. "
+          "Please set train_data_columns or eval_data_columns accordingly."
+      )
+
     parsed = {}
     for col in self.data_columns:
-      if col in features:
-        f = features[col]
-        if self.tokenize:
-          parsed[col] = np.array(f.bytes_list.value, dtype=object)
-        else:
-          parsed[col] = np.array(f.int64_list.value, dtype=np.int32)
+      f = features[col]
+      if self.tokenize:
+        parsed[col] = np.array(f.bytes_list.value, dtype=object)
+      else:
+        parsed[col] = np.array(f.int64_list.value, dtype=np.int32)
     return parsed
 
 
@@ -601,6 +607,12 @@ def __init__(self, feature_names: list[str]):
 
   def map(self, element: dict[str, Any]) -> dict[str, Any]:
     """Applies the feature filtering to the input element."""
+    missing = [n for n in self.feature_names if n not in element]
+    if missing:
+      raise ValueError(
+          f"Column {missing} not found in dataset. Available columns: {sorted(element.keys())}. "
+          "Please set train_data_columns or eval_data_columns accordingly."
+      )
     return {k: v for k, v in element.items() if k in self.feature_names}
 
 

src/maxtext/input_pipeline/tfds_data_processing.py

@@ -95,6 +95,13 @@ def preprocessing_pipeline(
     hf_access_token: str = "",
 ):
   """pipeline for preprocessing TFDS dataset."""
+  missing = [c for c in data_column_names if c not in dataset.element_spec]
+  if missing:
+    raise ValueError(
+        f"Column {missing} not found in dataset. Available columns: {sorted(dataset.element_spec.keys())}. "
+        "Please set train_data_columns or eval_data_columns accordingly."
+    )
+
   if not use_dpo:
     assert len(data_column_names) == 1
     dataset = dataset.map(