Enable zero-copy transfer of pyarrow buffer to ArrowDataset kernel when in eager mode (tensorflow#413)

Author: BryanCutler

tensorflow_io/arrow/kernels/arrow_dataset_ops.cc

@@ -583,11 +583,155 @@ class ArrowOpKernelBase : public DatasetOpKernel {
   std::vector<PartialTensorShape> output_shapes_;
 };
 
-// Op to create an ArrowDataset that consumes Arrow record batches from
-// memory in a Python process, or a Pandas DataFrame.
-class ArrowDatasetOp : public ArrowOpKernelBase {
+// Op to create an ArrowZeroCopyDataset that consumes Arrow record batches
+// from a memory buffer address owned in Python.
+class ArrowZeroCopyDatasetOp : public ArrowOpKernelBase {
  public:
-  explicit ArrowDatasetOp(OpKernelConstruction* ctx) : ArrowOpKernelBase(ctx) {}
+  explicit ArrowZeroCopyDatasetOp(OpKernelConstruction* ctx)
+      : ArrowOpKernelBase(ctx) {}
+
+  virtual void MakeArrowDataset(
+      OpKernelContext* ctx,
+      const std::vector<int32>& columns,
+      const int64 batch_size,
+      const ArrowBatchMode batch_mode,
+      const DataTypeVector& output_types,
+      const std::vector<PartialTensorShape>& output_shapes,
+      ArrowDatasetBase** output) override {
+    uintptr_t buffer_address;
+    OP_REQUIRES_OK(
+        ctx,
+        ParseScalarArgument<uintptr_t>(ctx, "buffer_address", &buffer_address));
+    const uint8_t* buffer = reinterpret_cast<const uint8_t*>(buffer_address);
+
+    int64_t buffer_size;
+    OP_REQUIRES_OK(
+        ctx,
+        ParseScalarArgument<int64_t>(ctx, "buffer_size", &buffer_size));
+    *output = new Dataset(
+        ctx,
+        buffer,
+        buffer_size,
+        columns,
+        batch_size,
+        batch_mode,
+        output_types_,
+        output_shapes_);
+  }
+
+ private:
+  class Dataset : public ArrowDatasetBase {
+   public:
+    Dataset(OpKernelContext* ctx,
+            const uint8_t* buffer_ptr,
+            const int64 buffer_size,
+            const std::vector<int32>& columns,
+            const int64 batch_size,
+            const ArrowBatchMode batch_mode,
+            const DataTypeVector& output_types,
+            const std::vector<PartialTensorShape>& output_shapes)
+        : ArrowDatasetBase(ctx, columns, batch_size, batch_mode,
+              output_types, output_shapes),
+          buffer_ptr_(buffer_ptr), buffer_size_(buffer_size) {}
+
+    string DebugString() const override {
+      return "ArrowZeroCopyDatasetOp::Dataset";
+    }
+
+   protected:
+    Status AsGraphDefInternal(SerializationContext* ctx,
+                              DatasetGraphDefBuilder* b,
+                              Node** output) const override {
+      Node* buffer = nullptr;
+      uintptr_t buffer_temp = reinterpret_cast<uintptr_t>(buffer_ptr_);
+      uint64 buffer_address = buffer_temp;
+      TF_RETURN_IF_ERROR(b->AddScalar(buffer_address, &buffer));
+      Node* size = nullptr;
+      TF_RETURN_IF_ERROR(
+          b->AddScalar(static_cast<int64>(buffer_size_), &size));
+      Node* columns = nullptr;
+      TF_RETURN_IF_ERROR(b->AddVector(columns_, &columns));
+      Node* batch_size = nullptr;
+      TF_RETURN_IF_ERROR(b->AddScalar(batch_size_, &batch_size));
+      Node* batch_mode = nullptr;
+      string batch_mode_str;
+      TF_RETURN_IF_ERROR(GetBatchModeStr(batch_mode_, &batch_mode_str));
+      TF_RETURN_IF_ERROR(b->AddScalar(batch_mode_str, &batch_mode));
+      TF_RETURN_IF_ERROR(
+          b->AddDataset(
+            this,
+            {buffer, size, columns, batch_size, batch_mode},
+            output));
+      return Status::OK();
+    }
+
+    std::unique_ptr<IteratorBase> MakeIteratorInternal(
+        const string& prefix) const override {
+      return std::unique_ptr<IteratorBase>(
+          new Iterator({this, strings::StrCat(prefix, "::Arrow")}));
+    }
+
+   private:
+    class Iterator : public ArrowBaseIterator<Dataset> {
+     public:
+      explicit Iterator(const Params& params)
+          : ArrowBaseIterator<Dataset>(params) {}
+
+     private:
+      Status SetupStreamsLocked() EXCLUSIVE_LOCKS_REQUIRED(mu_) override {
+        buffer_ = std::make_shared<arrow::Buffer>(
+            dataset()->buffer_ptr_,
+            dataset()->buffer_size_);
+        buffer_reader_ = std::make_shared<arrow::io::BufferReader>(buffer_);
+        CHECK_ARROW(
+            arrow::ipc::RecordBatchFileReader::Open(
+              buffer_reader_.get(),
+              buffer_->size(),
+              &reader_));
+        num_batches_ = reader_->num_record_batches();
+        if (num_batches_ > 0) {
+          CHECK_ARROW(
+              reader_->ReadRecordBatch(current_batch_idx_, &current_batch_));
+          TF_RETURN_IF_ERROR(CheckBatchColumnTypes(current_batch_));
+        }
+        return Status::OK();
+      }
+
+      Status NextStreamLocked() EXCLUSIVE_LOCKS_REQUIRED(mu_) override {
+        ArrowBaseIterator<Dataset>::NextStreamLocked();
+        if (++current_batch_idx_ < num_batches_) {
+          CHECK_ARROW(
+              reader_->ReadRecordBatch(current_batch_idx_, &current_batch_));
+        }
+        return Status::OK();
+      }
+
+      void ResetStreamsLocked() EXCLUSIVE_LOCKS_REQUIRED(mu_) override {
+        ArrowBaseIterator<Dataset>::ResetStreamsLocked();
+        reader_.reset();
+        current_batch_idx_ = 0;
+        num_batches_ = 0;
+      }
+
+      std::shared_ptr<arrow::Buffer> buffer_ GUARDED_BY(mu_);
+      std::shared_ptr<arrow::io::BufferReader> buffer_reader_ GUARDED_BY(mu_);
+      std::shared_ptr<arrow::ipc::RecordBatchFileReader> reader_
+          GUARDED_BY(mu_);
+      int current_batch_idx_ GUARDED_BY(mu_) = 0;
+      int num_batches_ GUARDED_BY(mu_) = 0;
+    };
+
+    const uint8_t* buffer_ptr_;
+    const int64 buffer_size_;
+  };
+};
+
+// Op to create an ArrowSerializedDataset that consumes Arrow record batches
+// serialized in a Tensor buffer.
+class ArrowSerializedDatasetOp : public ArrowOpKernelBase {
+ public:
+  explicit ArrowSerializedDatasetOp(OpKernelConstruction* ctx)
+      : ArrowOpKernelBase(ctx) {}
 
   virtual void MakeArrowDataset(
       OpKernelContext* ctx,
@@ -629,7 +773,9 @@ class ArrowDatasetOp : public ArrowOpKernelBase {
           batches_(std::move(batches_tensor)) {
     }
 
-    string DebugString() const override { return "ArrowDatasetOp::Dataset"; }
+    string DebugString() const override {
+      return "ArrowSerializedDatasetOp::Dataset";
+    }
 
    protected:
     Status AsGraphDefInternal(SerializationContext* ctx,
@@ -1009,8 +1155,11 @@ class ArrowStreamDatasetOp : public ArrowOpKernelBase {
   };
 };
 
-REGISTER_KERNEL_BUILDER(Name("ArrowDataset").Device(DEVICE_CPU),
-                        ArrowDatasetOp);
+REGISTER_KERNEL_BUILDER(Name("ArrowZeroCopyDataset").Device(DEVICE_CPU),
+                        ArrowZeroCopyDatasetOp);
+
+REGISTER_KERNEL_BUILDER(Name("ArrowSerializedDataset").Device(DEVICE_CPU),
+                        ArrowSerializedDatasetOp);
 
 REGISTER_KERNEL_BUILDER(Name("ArrowFeatherDataset").Device(DEVICE_CPU),
                         ArrowFeatherDatasetOp);

tensorflow_io/arrow/ops/dataset_ops.cc

@@ -19,7 +19,26 @@ limitations under the License.
 
 namespace tensorflow {
 
-REGISTER_OP("ArrowDataset")
+REGISTER_OP("ArrowZeroCopyDataset")
+    .Input("buffer_address: uint64")
+    .Input("buffer_size: int64")
+    .Input("columns: int32")
+    .Input("batch_size: int64")
+    .Input("batch_mode: string")
+    .Output("handle: variant")
+    .Attr("output_types: list(type) >= 1")
+    .Attr("output_shapes: list(shape) >= 1")
+    .SetIsStateful()
+    .SetShapeFn(shape_inference::ScalarShape)
+    .Doc(R"doc(
+Creates a dataset that zero-copy reads data from an Arrow Buffer.
+
+buffer_address: Buffer address as long int with contents as Arrow RecordBatches
+in file format.
+buffer_size: Buffer size in bytes
+)doc");
+
+REGISTER_OP("ArrowSerializedDataset")
     .Input("serialized_batches: string")
     .Input("columns: int32")
     .Input("batch_size: int64")

tensorflow_io/arrow/python/ops/arrow_dataset_ops.py

@@ -161,13 +161,14 @@ def __init__(self,
                output_types,
                output_shapes=None,
                batch_size=None,
-               batch_mode='keep_remainder'):
+               batch_mode='keep_remainder',
+               arrow_buffer=None):
     """Create an ArrowDataset from a Tensor of serialized batches.
     This constructor requires pyarrow to be installed.
 
     Args:
       serialized_batches: A string Tensor as a serialized buffer containing
-                          Arrow record batches as Arrow file format
+                          Arrow record batches in Arrow File format
       columns: A list of column indices to be used in the Dataset
       output_types: Tensor dtypes of the output tensors
       output_shapes: TensorShapes of the output tensors or None to
@@ -180,9 +181,39 @@ def __init__(self,
                   "keep_remainder" (default, keeps partial batch data),
                   "drop_remainder" (discard partial batch data),
                   "auto" (size to number of records in Arrow record batch)
+      arrow_buffer: Optional Arrow Buffer containing Arrow record batches in
+                    Arrow File format. This will share the Arrow buffer with
+                    the C++ kernel by address for zero-copy. Only supported if
+                    the kernel process is local, with TensorFlow in eager mode.
+                    If this is used, set `serialized_batches` to `None`.
     """
+    if serialized_batches is not None:
+      make_variant_fn = partial(
+          core_ops.arrow_serialized_dataset,
+          serialized_batches)
+    elif arrow_buffer is None:
+      raise ValueError("Must set either serialzied_batches or arrow_buffer")
+    elif not tf.executing_eagerly():
+      raise ValueError("Using arrow_buffer for zero-copy only supported in "
+                       "TensorFlow Eager mode.")
+    else:
+      address_int = arrow_buffer.address
+      buffer_address = tf.convert_to_tensor(
+          address_int,
+          dtype=dtypes.uint64,
+          name="buffer_address")
+      buffer_size = tf.convert_to_tensor(
+          arrow_buffer.size,
+          dtype=dtypes.int64,
+          name="buffer_size")
+      make_variant_fn = partial(
+          core_ops.arrow_zero_copy_dataset,
+          buffer_address,
+          buffer_size)
+      # Keep a reference to the arrow buffers used
+      self._arrow_buffer_refs = [arrow_buffer]
     super(ArrowDataset, self).__init__(
-        partial(core_ops.arrow_dataset, serialized_batches),
+        make_variant_fn,
         columns,
         output_types,
         output_shapes,
@@ -221,22 +252,33 @@ def from_record_batches(cls,
     if columns is None:
       columns = tuple(range(record_batches[0].num_columns))
     assert record_batches
-    buf = io.BytesIO()
-    writer = pa.RecordBatchFileWriter(buf, record_batches[0].schema)
-    for batch in record_batches:
-      writer.write_batch(batch)
-    writer.close()
-    serialized_batches = tf.convert_to_tensor(
-        buf.getvalue(),
-        dtype=dtypes.string,
-        name="serialized_batches")
+    if tf.executing_eagerly():
+      sink = pa.BufferOutputStream()
+      writer = pa.RecordBatchFileWriter(sink, record_batches[0].schema)
+      for batch in record_batches:
+        writer.write_batch(batch)
+      writer.close()
+      serialized_batches = None
+      arrow_buffer = sink.getvalue()
+    else:
+      buf = io.BytesIO()
+      writer = pa.RecordBatchFileWriter(buf, record_batches[0].schema)
+      for batch in record_batches:
+        writer.write_batch(batch)
+      writer.close()
+      serialized_batches = tf.convert_to_tensor(
+          buf.getvalue(),
+          dtype=dtypes.string,
+          name="serialized_batches")
+      arrow_buffer = None
     return cls(
         serialized_batches,
         columns,
         output_types,
         output_shapes,
-        batch_size,
-        batch_mode)
+        batch_size=batch_size,
+        batch_mode=batch_mode,
+        arrow_buffer=arrow_buffer)
 
   @classmethod
   def from_pandas(cls,