Merge remote-tracking branch 'upstream/master' into core

Author: hgarrereyn

Diff for: RELEASE.md

@@ -17,6 +17,7 @@
     check and prevent duplications. DO NOT rely on this behavior.
 *   CsvTFXIO now allows skipping CSV headers (`set skip_header_lines`).
 *   CsvTFXIO now requires `telemetry_descriptors` to construct.
+*   Depends on `pyarrow>=0.17,<0.18`.
 
 ## Breaking changes
 

Diff for: tfx_bsl/beam/run_inference.py

@@ -545,6 +545,8 @@ def setup(self):
   def _setup_model(
       self, inference_spec_type: model_spec_pb2.InferenceSpecType
   ):
+    self._ai_platform_prediction_model_spec = (
+        inference_spec_type.ai_platform_prediction_model_spec)
     self._api_client = None
 
     project_id = (
@@ -589,10 +591,10 @@ def _make_request(self, body: Mapping[Text, List[Any]]) -> http.HttpRequest:
     return self._api_client.projects().predict(
         name=self._full_model_name, body=body)
 
-  @classmethod
-  def _prepare_instances(
-      cls, elements: List[ExampleType]
+  def _prepare_instances_dict(
+      self, elements: List[ExampleType]
   ) -> Generator[Mapping[Text, Any], None, None]:
+    """Prepare instances by converting features to dictionary."""
     for example in elements:
       # TODO(b/151468119): support tf.train.SequenceExample
       if not isinstance(example, tf.train.Example):
@@ -604,13 +606,28 @@ def _prepare_instances(
         if attr_name is None:
           continue
         attr = getattr(feature, attr_name)
-        values = cls._parse_feature_content(attr.value, attr_name,
-                                            cls._sending_as_binary(input_name))
+        values = self._parse_feature_content(
+            attr.value, attr_name, self._sending_as_binary(input_name))
         # Flatten a sequence if its length is 1
         values = (values[0] if len(values) == 1 else values)
         instance[input_name] = values
       yield instance
 
+  def _prepare_instances_serialized(
+      self, elements: List[ExampleType]
+  ) -> Generator[Mapping[Text, Text], None, None]:
+    """Prepare instances by base64 encoding serialized examples."""
+    for example in elements:
+      yield {'b64': base64.b64encode(example.SerializeToString()).decode()}
+
+  def _prepare_instances(
+      self, elements: List[ExampleType]
+  ) -> Generator[Mapping[Text, Any], None, None]:
+    if self._ai_platform_prediction_model_spec.use_serialization_config:
+      return self._prepare_instances_serialized(elements)
+    else:
+      return self._prepare_instances_dict(elements)
+
   @staticmethod
   def _sending_as_binary(input_name: Text) -> bool:
     """Whether data should be sent as binary."""

Diff for: tfx_bsl/beam/run_inference_test.py

@@ -18,6 +18,7 @@
 # Standard __future__ imports
 from __future__ import print_function
 
+import base64
 import json
 import os
 try:
@@ -603,9 +604,17 @@ def test_request_body_with_binary_data(self):
         feature { key: "z" value { float_list { value: [4.5, 5, 5.5] }}}
       }
       """, tf.train.Example())
-    result = list(
-        run_inference._RemotePredictDoFn._prepare_instances([example]))
-    self.assertEqual([
+    inference_spec_type = model_spec_pb2.InferenceSpecType(
+        ai_platform_prediction_model_spec=model_spec_pb2
+        .AIPlatformPredictionModelSpec(
+            project_id='test_project',
+            model_name='test_model',
+            version_name='test_version'))
+    remote_predict = run_inference._RemotePredictDoFn(
+      None, fixed_inference_spec_type=inference_spec_type)
+    remote_predict._setup_model(remote_predict._fixed_inference_spec_type)
+    result = list(remote_predict._prepare_instances([example]))
+    self.assertEqual(result, [
         {
             'x_bytes': {
                 'b64': 'QVNhOGFzZGY='
@@ -614,7 +623,31 @@ def test_request_body_with_binary_data(self):
             'y': [1, 2],
             'z': [4.5, 5, 5.5]
         },
-    ], result)
+    ])
+
+  def test_request_serialized_example(self):
+    example = text_format.Parse(
+        """
+      features {
+        feature { key: "x_bytes" value { bytes_list { value: ["ASa8asdf"] }}}
+        feature { key: "x" value { bytes_list { value: "JLK7ljk3" }}}
+        feature { key: "y" value { int64_list { value: [1, 2] }}}
+      }
+      """, tf.train.Example())
+    inference_spec_type = model_spec_pb2.InferenceSpecType(
+        ai_platform_prediction_model_spec=model_spec_pb2
+        .AIPlatformPredictionModelSpec(
+            project_id='test_project',
+            model_name='test_model',
+            version_name='test_version',
+            use_serialization_config=True))
+    remote_predict = run_inference._RemotePredictDoFn(
+      None, fixed_inference_spec_type=inference_spec_type)
+    remote_predict._setup_model(remote_predict._fixed_inference_spec_type)
+    result = list(remote_predict._prepare_instances([example]))
+    self.assertEqual(result, [{
+        'b64': base64.b64encode(example.SerializeToString()).decode()
+    }])
 
 
 class RunInferenceCoreTest(RunInferenceFixture):

Diff for: tfx_bsl/cc/sketches/BUILD

@@ -40,26 +40,6 @@ cc_test(
     ],
 )
 
-cc_library(
-    name = "sketches_submodule",
-    srcs = ["sketches_submodule.cc"],
-    hdrs = ["sketches_submodule.h"],
-    copts = [
-        "-fexceptions",
-    ],
-    features = ["-use_header_modules"],
-    visibility = [
-        "//tfx_bsl/cc:__pkg__",
-    ],
-    deps = [
-        ":kmv_sketch",
-        "//tfx_bsl/cc/pybind11:absl_casters",
-        "//tfx_bsl/cc/pybind11:arrow_casters",
-        "@com_google_absl//absl/strings",
-        "@pybind11",
-    ],
-)
-
 cc_library(
     name = "misragries_sketch",
     srcs = ["misragries_sketch.cc"],
@@ -88,3 +68,24 @@ cc_test(
         "@com_google_googletest//:gtest_main",
     ],
 )
+
+cc_library(
+    name = "sketches_submodule",
+    srcs = ["sketches_submodule.cc"],
+    hdrs = ["sketches_submodule.h"],
+    copts = [
+        "-fexceptions",
+    ],
+    features = ["-use_header_modules"],
+    visibility = [
+        "//tfx_bsl/cc:__pkg__",
+    ],
+    deps = [
+        ":kmv_sketch",
+        ":misragries_sketch",
+        "//tfx_bsl/cc/pybind11:absl_casters",
+        "//tfx_bsl/cc/pybind11:arrow_casters",
+        "@com_google_absl//absl/strings",
+        "@pybind11",
+    ],
+)

Diff for: tfx_bsl/cc/sketches/kmv_sketch.h

@@ -46,17 +46,16 @@ class KmvSketch {
   // KmvSketch is copyable and movable.
   KmvSketch(int num_buckets);
   ~KmvSketch() = default;
-  // Updates the sketch with an Arrow array of values. Supports numeric arrays
-  // of all integral types, binary arrays, and string arrays.
+  // Updates the sketch with an Arrow array of values.
   Status AddValues(const arrow::Array& array);
-  // Merges the sketch from another object into this sketch. Returns error if
-  // the other sketch has a different number of buckets than this sketch.
+  // Merges another KMV sketch into this sketch. Returns error if the other
+  // sketch has a different number of buckets than this sketch.
   Status Merge(KmvSketch& other);
   // Estimates the number of distinct elements.
   uint64_t Estimate() const;
   // Serializes the sketch into a string.
   std::string Serialize() const;
-  // Converts the encoded sketch into a KmvSketch object.
+  // Deserializes the string to a KmvSketch object.
   static KmvSketch Deserialize(absl::string_view encoded);
 
  private:

Diff for: tfx_bsl/cc/sketches/misragries_sketch.cc

@@ -131,7 +131,8 @@ class UpdateItemCountsVisitor : public arrow::ArrayVisitor  {
         if (array.IsNull(i)) {
           continue;
         }
-        const auto item = array.GetView(i);
+        const auto value = array.GetView(i);
+        const auto item = absl::string_view(value.data(), value.size());
         const float weight = weights_->Value(i);
         InsertItem(item, weight);
       }

Diff for: tfx_bsl/cc/sketches/misragries_sketch.h

@@ -65,10 +65,10 @@ class MisraGriesSketch {
   MisraGriesSketch(int num_buckets);
   // This class is copyable and movable.
   ~MisraGriesSketch() = default;
-  // Add an array of items.
+  // Adds an array of items.
   Status AddValues(const arrow::Array& items);
-  // Add an array of items with their associated weights. Raises an error if the
-  // weights are not a FloatArray.
+  // Adds an array of items with their associated weights. Raises an error if
+  // the weights are not a FloatArray.
   Status AddValues(const arrow::Array& items, const arrow::Array& weights);
   // Merges another MisraGriesSketch into this sketch. Raises an error if the
   // sketches do not have the same number of buckets.
@@ -77,9 +77,9 @@ class MisraGriesSketch {
   std::vector<std::pair<std::string, double>> GetCounts() const;
   // Creates a struct array <values, counts> of the top-k items.
   Status Estimate(std::shared_ptr<arrow::Array>* values_and_counts_array) const;
-  // Serializes proto of sketch.
+  // Serializes the sketch into a string.
   std::string Serialize() const;
-  // Parses encoded MisraGries proto.
+  // Deserializes the string to a MisraGries object.
   static MisraGriesSketch Deserialize(absl::string_view encoded);
   // Gets delta_.
   double GetDelta() const;

Diff for: tfx_bsl/cc/sketches/misragries_sketch_test.cc

@@ -387,7 +387,7 @@ TEST(MisraGriesSketchTest, AddWeightedZipfDistribution) {
     for (auto const& item : counts) {
       double true_count = true_counts[std::string(item.first)] * trial;
       double estimated_count = item.second;
-      EXPECT_LE(true_count, estimated_count + 1e8);
+      EXPECT_LE(true_count, estimated_count + 1e-8);
       EXPECT_LE(estimated_count - mg.GetDelta(), true_count);
     }
   }