Preemptive fixes for upcoming Beam PR 10717, which will enable Python 3 annotations support by default.

PiperOrigin-RevId: 299228970

Author: tfx-copybara

tensorflow_data_validation/api/stats_api.py

@@ -48,7 +48,6 @@
 import apache_beam as beam
 import pyarrow as pa
 from tensorflow_data_validation import constants
-from tensorflow_data_validation import types
 from tensorflow_data_validation.statistics import stats_impl
 from tensorflow_data_validation.statistics import stats_options
 from typing import Generator
@@ -122,8 +121,8 @@ def expand(self, dataset: beam.pvalue.PCollection) -> beam.pvalue.PCollection:
             stats_impl.GenerateStatisticsImpl(self._options))
 
 
-def _sample_at_rate(example: types.Example, sample_rate: float
-                   ) -> Generator[types.Example, None, None]:
+def _sample_at_rate(example: pa.Table, sample_rate: float
+                   ) -> Generator[pa.Table, None, None]:
   """Sample examples at input sampling rate."""
   # TODO(pachristopher): Revisit this to decide if we need to fix a seed
   # or add an optional seed argument.

tensorflow_data_validation/statistics/generators/lift_stats_generator.py

@@ -185,7 +185,7 @@ def _get_example_value_presence(
 def _to_partial_copresence_counts(
     sliced_table: types.SlicedTable, y_path: types.FeaturePath,
     x_paths: Iterable[types.FeaturePath],
-    y_boundaries: Optional[Iterable[float]], weight_column_name: Optional[Text]
+    y_boundaries: Optional[np.ndarray], weight_column_name: Optional[Text]
 ) -> Iterator[Tuple[_SlicedXYKey, _CountType]]:
   """Yields per-(slice, path_x, x, y) counts of examples with x and y.
 
@@ -234,7 +234,7 @@ def _to_partial_copresence_counts(
 
 def _to_partial_counts(
     sliced_table: types.SlicedTable, path: types.FeaturePath,
-    boundaries: Optional[Iterable[float]], weight_column_name: Optional[Text]
+    boundaries: Optional[np.ndarray], weight_column_name: Optional[Text]
 ) -> Iterator[Tuple[Tuple[types.SliceKey, Union[_XType, _YType]], _CountType]]:
   """Yields per-(slice, value) counts of the examples with value in path."""
   slice_key, table = sliced_table
@@ -274,9 +274,9 @@ def _get_unicode_value(value: Union[Text, bytes], path: types.FeaturePath):
 
 
 def _make_dataset_feature_stats_proto(
-    lifts: Tuple[_SlicedFeatureKey, _LiftSeries], y_path: types.FeaturePath,
-    y_boundaries: Optional[np.ndarray], weighted_examples: bool,
-    output_custom_stats: bool
+    lifts: Tuple[_SlicedFeatureKey, Iterable[_LiftSeries]],
+    y_path: types.FeaturePath, y_boundaries: Optional[np.ndarray],
+    weighted_examples: bool, output_custom_stats: bool
 ) -> Tuple[types.SliceKey, statistics_pb2.DatasetFeatureStatistics]:
   """Generates DatasetFeatureStatistics proto for a given x_path, y_path pair.
 

tensorflow_data_validation/statistics/generators/top_k_uniques_stats_generator.py

@@ -155,7 +155,7 @@ def _make_dataset_feature_stats_proto_with_topk_for_single_feature(
                                                   FeaturePathTuple],
                                             List[FeatureValueCount]],
     categorical_features: Set[types.FeaturePath], is_weighted_stats: bool,
-    num_top_values: int, frequency_threshold: float,
+    num_top_values: int, frequency_threshold: Union[int, float],
     num_rank_histogram_buckets: int
 ) -> Tuple[types.SliceKey, statistics_pb2.DatasetFeatureStatistics]:
   """Makes a DatasetFeatureStatistics proto with top-k stats for a feature."""
@@ -198,12 +198,12 @@ def _weighted_unique(values: np.ndarray, weights: np.ndarray
 
 
 def _to_topk_tuples(
-    sliced_table: Tuple[Text, pa.Table],
+    sliced_table: Tuple[types.SliceKey, pa.Table],
     bytes_features: FrozenSet[types.FeaturePath],
     categorical_features: FrozenSet[types.FeaturePath],
     weight_feature: Optional[Text]
 ) -> Iterable[
-    Tuple[Tuple[Text, FeaturePathTuple, Any],
+    Tuple[Tuple[types.SliceKey, FeaturePathTuple, Any],
           Union[int, Tuple[int, Union[int, float]]]]]:
   """Generates tuples for computing top-k and uniques from input tables."""
   slice_key, table = sliced_table

tensorflow_data_validation/statistics/stats_impl.py

@@ -118,7 +118,7 @@ def __init__(
     """
     self._options = options
     self._is_slicing_enabled = (
-        is_slicing_enabled or self._options.slice_functions)
+        is_slicing_enabled or bool(self._options.slice_functions))
 
   def expand(self, dataset: beam.pvalue.PCollection) -> beam.pvalue.PCollection:
     # Handles generators by their type: