Adding cord support for ZlibOutputBuffer. This improves snapshot benchmark performance by ~2%.

PiperOrigin-RevId: 257679909

tensorflow/core/kernels/data/experimental/snapshot_dataset_op.cc

@@ -91,12 +91,7 @@ class SnapshotWriter {
 
     TF_RETURN_IF_ERROR(dest_->Append(StringPiece(header, sizeof(header))));
 
-    // TODO(frankchn): Remove after ZlibOutputBuffer Cord support is added.
-    if (compression_type_.empty()) {
-      return dest_->Append(data);
-    } else {
-      return dest_->Append(data.ToString());
-    }
+    return dest_->Append(data);
   }
 #endif  // PLATFORM_GOOGLE
 

tensorflow/core/lib/io/zlib_outputbuffer.cc

@@ -190,6 +190,17 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
   return Status::OK();
 }
 
+#if defined(PLATFORM_GOOGLE)
+Status ZlibOutputBuffer::Append(const absl::Cord& cord) {
+  absl::CordReader reader(cord);
+  absl::string_view fragment;
+  while (reader.ReadFragment(&fragment)) {
+    TF_RETURN_IF_ERROR(Append(fragment));
+  }
+  return Status::OK();
+}
+#endif
+
 Status ZlibOutputBuffer::Flush() {
   TF_RETURN_IF_ERROR(DeflateBuffered(Z_PARTIAL_FLUSH));
   TF_RETURN_IF_ERROR(FlushOutputBufferToFile());

tensorflow/core/lib/io/zlib_outputbuffer.h

@@ -65,6 +65,10 @@ class ZlibOutputBuffer : public WritableFile {
   // To immediately write contents to file call `Flush()`.
   Status Append(StringPiece data) override;
 
+#if defined(PLATFORM_GOOGLE)
+  Status Append(const absl::Cord& cord) override;
+#endif
+
   // Deflates any cached input and writes all output to file.
   Status Flush() override;