Merge pull request #51189 from tensorflow/mm-cherrypick-e0b6e58c32805…

…9829c3eb968136f17aa72b6c876-on-r2.4

Fix segfault/heap buffer overflow in `{Experimental,}DatasetToTFRecor…

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

@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/core/framework/function_handle_cache.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/data/dataset_utils.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/threadpool.h"
@@ -87,8 +88,20 @@ class ToTFRecordOp : public AsyncOpKernel {
     TF_RETURN_IF_ERROR(dataset->MakeIterator(
         &iter_ctx, /*parent=*/nullptr, "ToTFRecordOpIterator", &iterator));
 
+    const int num_output_dtypes = dataset->output_dtypes().size();
+    if (num_output_dtypes != 1) {
+      return errors::InvalidArgument(
+          "ToTFRecordOp currently only support datasets of 1 single column, ",
+          "but got ", num_output_dtypes);
+    }
+    const DataType dt = dataset->output_dtypes()[0];
+    if (dt != DT_STRING) {
+      return errors::InvalidArgument(
+          "ToTFRecordOp currently only supports DT_STRING dataypes, but got ",
+          DataTypeString(dt));
+    }
     std::vector<Tensor> components;
-    components.reserve(dataset->output_dtypes().size());
+    components.reserve(num_output_dtypes);
     bool end_of_sequence;
     do {
       TF_RETURN_IF_ERROR(