Update get_weight_feature to be used with Arrow table inputs and use same in arrow_util.enumerate_arrays.

PiperOrigin-RevId: 272281600

Author: caveness

tensorflow_data_validation/arrow/arrow_util.py

@@ -35,6 +35,45 @@
     pywrap.TFDV_Arrow_MakeListArrayFromParentIndicesAndValues)
 
 
+def _get_weight_feature(input_table: pa.Table,
+                        weight_feature: Text) -> np.ndarray:
+  """Gets the weight column from the input table.
+
+  Args:
+    input_table: Input table.
+    weight_feature: Name of the weight feature.
+
+  Returns:
+    A numpy array containing the weights of the examples in the input table.
+
+  Raises:
+    ValueError: If the weight feature is not present in the input table or is
+        not a valid weight feature (must be of numeric type and have a
+        single value for each example).
+  """
+  try:
+    weights = input_table.column(weight_feature).data.chunk(0)
+  except KeyError:
+    raise ValueError('Weight feature "{}" not present in the input '
+                     'table.'.format(weight_feature))
+
+  # Before flattening, check that there is a single value for each example.
+  weight_lengths = ListLengthsFromListArray(weights).to_numpy()
+  if not np.all(weight_lengths == 1):
+    raise ValueError(
+        'Weight feature "{}" must have exactly one value in each example.'
+        .format(weight_feature))
+  weights = weights.flatten()
+  # Before converting to numpy view, check the type (cannot convert string and
+  # binary arrays to numpy view).
+  weights_type = weights.type
+  if pa.types.is_string(weights_type) or pa.types.is_binary(weights_type):
+    raise ValueError(
+        'Weight feature "{}" must be of numeric type. Found {}.'.format(
+            weight_feature, weights_type))
+  return weights.to_numpy()
+
+
 def primitive_array_to_numpy(primitive_array: pa.Array) -> np.ndarray:
   """Converts a primitive Arrow array to a numpy 1-D ndarray.
 
@@ -122,10 +161,7 @@ def _recursion_helper(
 
   weights = None
   if weight_column is not None:
-    weights = table.column(weight_column).data.chunk(0).flatten().to_numpy()
-    if weights.size != table.num_rows:
-      raise ValueError(
-          'The weight feature must have exactly one value in each example')
+    weights = _get_weight_feature(table, weight_column)
   for column in table.columns:
     column_name = column.name
     # use "yield from" after PY 3.3.

tensorflow_data_validation/arrow/arrow_util_test.py

@@ -236,17 +236,39 @@ def testMakeListArray(self, num_parents, parent_indices, values, expected):
 
 class EnumerateArraysTest(absltest.TestCase):
 
-  def testInvalidWeightColumn(self):
+  def testInvalidWeightColumnMissingValue(self):
     with self.assertRaisesRegex(
         ValueError,
-        "weight feature must have exactly one value in each example"):
+        'Weight feature "w" must have exactly one value.*'):
       for _ in arrow_util.enumerate_arrays(
           pa.Table.from_arrays([pa.array([[1], [2, 3]]),
                                 pa.array([[1], []])], ["v", "w"]),
           weight_column="w",
           enumerate_leaves_only=False):
         pass
 
+  def testInvalidWeightColumnTooManyValues(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        'Weight feature "w" must have exactly one value.*'):
+      for _ in arrow_util.enumerate_arrays(
+          pa.Table.from_arrays([pa.array([[1], [2, 3]]),
+                                pa.array([[1], [2, 2]])], ["v", "w"]),
+          weight_column="w",
+          enumerate_leaves_only=False):
+        pass
+
+  def testInvalidWeightColumnStringValues(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        'Weight feature "w" must be of numeric type.*'):
+      for _ in arrow_util.enumerate_arrays(
+          pa.Table.from_arrays([pa.array([[1], [2, 3]]),
+                                pa.array([["two"], ["two"]])], ["v", "w"]),
+          weight_column="w",
+          enumerate_leaves_only=False):
+        pass
+
   def testEnumerate(self):
     input_table = pa.Table.from_arrays([
         pa.array([[1], [2, 3]]),

tensorflow_data_validation/utils/stats_util.py

@@ -21,7 +21,7 @@
 import numpy as np
 from tensorflow_data_validation import types
 from tensorflow_data_validation.pyarrow_tf import pyarrow as pa
-from typing import Dict, List, Optional, Text, Union
+from typing import Dict, Optional, Text, Union
 from google.protobuf import text_format
 # TODO(b/125849585): Update to import from TF directly.
 from tensorflow.python.lib.io import file_io  # pylint: disable=g-direct-tensorflow-import
@@ -175,43 +175,6 @@ def _make_feature_stats_proto(
   return result
 
 
-def get_weight_feature(input_batch: types.ExampleBatch,
-                       weight_feature: types.FeatureName) -> List[np.ndarray]:
-  """Gets the weight feature from the input batch.
-
-  Args:
-    input_batch: Input batch of examples.
-    weight_feature: Name of the weight feature.
-
-  Returns:
-    A list containing the weights of the examples in the input batch.
-
-  Raises:
-    ValueError: If the weight feature is not present in the input batch or is
-        not a valid weight feature (must be of numeric type and have a
-        single value).
-  """
-  try:
-    weights = input_batch[weight_feature]
-  except KeyError:
-    raise ValueError('Weight feature "{}" not present in the input '
-                     'batch.'.format(weight_feature))
-
-  # Check if we have a valid weight feature.
-  for w in weights:
-    if w is None:
-      raise ValueError('Weight feature "{}" missing in an '
-                       'example.'.format(weight_feature))
-    elif (get_feature_type(w.dtype) ==
-          statistics_pb2.FeatureNameStatistics.STRING):
-      raise ValueError('Weight feature "{}" must be of numeric type. '
-                       'Found {}.'.format(weight_feature, w))
-    elif w.size != 1:
-      raise ValueError('Weight feature "{}" must have a single value. '
-                       'Found {}.'.format(weight_feature, w))
-  return weights  # pytype: disable=bad-return-type
-
-
 def write_stats_text(stats: statistics_pb2.DatasetFeatureStatisticsList,
                      output_path: bytes) -> None:
   """Writes a DatasetFeatureStatisticsList proto to a file in text format.

tensorflow_data_validation/utils/stats_util_test.py

@@ -122,33 +122,6 @@ def test_make_dataset_feature_stats_proto(self):
           expected[types.FeaturePath.from_proto(actual_feature_stats.path)],
           normalize_numbers=True)
 
-  def test_get_weight_feature_with_valid_weight_feature(self):
-    batch = {'a': np.array([np.array([1, 2]), np.array([3])]),
-             'w': np.array([np.array([10]), np.array([20])])}
-    actual = stats_util.get_weight_feature(batch, 'w')
-    np.testing.assert_equal(actual, batch['w'])
-
-  def test_get_weight_feature_invalid_weight_feature(self):
-    batch = {'a': np.array([np.array([1])])}
-    with self.assertRaisesRegexp(ValueError, 'Weight feature.*not present'):
-      stats_util.get_weight_feature(batch, 'w')
-
-  def test_get_weight_feature_with_weight_feature_missing(self):
-    batch = {'a': np.array([np.array([1])]), 'w': np.array([None])}
-    with self.assertRaisesRegexp(ValueError, 'Weight feature.*missing'):
-      stats_util.get_weight_feature(batch, 'w')
-
-  def test_get_weight_feature_with_weight_feature_string_type(self):
-    batch = {'a': np.array([np.array([1])]), 'w': np.array([np.array(['a'])])}
-    with self.assertRaisesRegexp(ValueError, 'Weight feature.*numeric type'):
-      stats_util.get_weight_feature(batch, 'w')
-
-  def test_get_weight_feature_with_weight_feature_multiple_values(self):
-    batch = {'a': np.array([np.array([1])]),
-             'w': np.array([np.array([2, 3])])}
-    with self.assertRaisesRegexp(ValueError, 'Weight feature.*single value'):
-      stats_util.get_weight_feature(batch, 'w')
-
   def test_get_utf8(self):
     self.assertEqual(u'This is valid.',
                      stats_util.maybe_get_utf8(b'This is valid.'))