Adds a helper to safely fetch a column from a RecordBatch and updates a few uses.

PiperOrigin-RevId: 397365531

Author: tfx-copybara

tensorflow_data_validation/arrow/arrow_util.py

@@ -342,3 +342,29 @@ def get_nest_level(array_type: pa.DataType) -> int:
   if pa.types.is_null(array_type):
     result += 1
   return result
+
+
+def get_column(record_batch: pa.RecordBatch,
+               feature_name: types.FeatureName,
+               missing_ok: bool = False) -> Optional[pa.Array]:
+  """Get a column by feature name.
+
+  Args:
+    record_batch: A pa.RecordBatch.
+    feature_name: The name of a feature (column) within record_batch.
+    missing_ok: If True, returns None for missing feature names.
+
+  Returns:
+    The column with the specified name, or None if missing_ok is true and
+    a column with the specified name is missing, or more than one exist.
+
+  Raises:
+    KeyError: If a column with the specified name is missing, or more than
+    one exist, and missing_ok is False.
+  """
+  idx = record_batch.schema.get_field_index(feature_name)
+  if idx < 0:
+    if missing_ok:
+      return None
+    raise KeyError('missing column %s' % feature_name)
+  return record_batch.column(idx)

tensorflow_data_validation/arrow/arrow_util_test.py

@@ -340,6 +340,14 @@ def testFlattenNestedNonList(self):
     self.assertTrue(flattened.equals(pa.array([1, 2])))
     np.testing.assert_array_equal(parent_indices, [0, 1])
 
+  def testGetColumn(self):
+    self.assertTrue(
+        arrow_util.get_column(_INPUT_RECORD_BATCH,
+                              "f1").equals(pa.array([[1], [2, 3]])))
+    self.assertIsNone(
+        arrow_util.get_column(_INPUT_RECORD_BATCH, "xyz", missing_ok=True))
+    with self.assertRaises(KeyError):
+      arrow_util.get_column(_INPUT_RECORD_BATCH, "xyz")
 
 if __name__ == "__main__":
   absltest.main()

tensorflow_data_validation/arrow/decoded_examples_to_arrow_test.py

@@ -23,6 +23,7 @@
 import numpy as np
 import pyarrow as pa
 import six
+from tensorflow_data_validation.arrow import arrow_util
 from tensorflow_data_validation.arrow import decoded_examples_to_arrow
 
 
@@ -189,8 +190,7 @@ def test_conversion(self, input_examples, expected_output):
         input_examples)
     self.assertLen(expected_output, record_batch.num_columns)
     for feature_name, expected_arrow_array in six.iteritems(expected_output):
-      actual = record_batch.column(
-          record_batch.schema.get_field_index(feature_name))
+      actual = arrow_util.get_column(record_batch, feature_name)
       self.assertTrue(
           expected_arrow_array.equals(actual),
           "{} vs {}".format(expected_arrow_array, actual))

tensorflow_data_validation/statistics/stats_impl.py

@@ -339,15 +339,14 @@ def _filter_features(
   Returns:
     An Arrow RecordBatch containing only features on the allowlist.
   """
-  schema = record_batch.schema
-  column_names = set(schema.names)
   columns_to_select = []
   column_names_to_select = []
   for feature_name in feature_allowlist:
-    if feature_name in column_names:
-      columns_to_select.append(
-          record_batch.column(schema.get_field_index(feature_name)))
-      column_names_to_select.append(feature_name)
+    col = arrow_util.get_column(record_batch, feature_name, missing_ok=True)
+    if col is None:
+      continue
+    columns_to_select.append(col)
+    column_names_to_select.append(feature_name)
   return pa.RecordBatch.from_arrays(columns_to_select, column_names_to_select)
 
 
@@ -523,8 +522,7 @@ def add_input(self, accumulator: List[float],
                 examples: pa.RecordBatch) -> List[float]:
     accumulator[0] += examples.num_rows
     if self._weight_feature:
-      weights_column = examples.column(
-          examples.schema.get_field_index(self._weight_feature))
+      weights_column = arrow_util.get_column(examples, self._weight_feature)
       accumulator[1] += np.sum(np.asarray(weights_column.flatten()))
     return accumulator
 
@@ -787,13 +785,13 @@ def generate_partial_statistics_in_memory(
   """
   result = []
   if options.feature_allowlist:
-    schema = record_batch.schema
-    columns = [
-        record_batch.column(schema.get_field_index(f))
-        for f in options.feature_allowlist
-    ]
-    record_batch = pa.RecordBatch.from_arrays(columns,
-                                              list(options.feature_allowlist))
+    columns, features = [], []
+    for feature_name in options.feature_allowlist:
+      c = arrow_util.get_column(record_batch, feature_name, missing_ok=True)
+      if c is not None:
+        columns.append(c)
+        features.append(feature_name)
+    record_batch = pa.RecordBatch.from_arrays(columns, features)
   for generator in stats_generators:
     result.append(
         generator.add_input(generator.create_accumulator(), record_batch))

tensorflow_data_validation/utils/slicing_util.py

@@ -123,12 +123,12 @@ def feature_value_slicer(record_batch: pa.RecordBatch) -> Iterable[
     """
     per_feature_parent_indices = []
     for feature_name, values in six.iteritems(features):
-      idx = record_batch.schema.get_field_index(feature_name)
+      feature_array = arrow_util.get_column(
+          record_batch, feature_name, missing_ok=True)
       # If the feature name does not appear in the schema for this record batch,
       # drop it from the set of sliced features.
-      if idx < 0:
+      if feature_array is None:
         continue
-      feature_array = record_batch.column(idx)
       flattened, value_parent_indices = arrow_util.flatten_nested(
           feature_array, True)
       non_missing_values = np.asarray(flattened)