Add arrow binary data type support (#1702)

* Add arrow binary data type support

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Add test case for string type

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

Author: yongtang

tensorflow_io/core/kernels/arrow/arrow_util.cc

@@ -33,6 +33,10 @@ Status GetTensorFlowType(std::shared_ptr<::arrow::DataType> dtype,
     *out = ::tensorflow::DT_STRING;
     return Status::OK();
   }
+  if (dtype->id() == ::arrow::Type::BINARY) {
+    *out = ::tensorflow::DT_STRING;
+    return Status::OK();
+  }
   ::arrow::Status status =
       ::arrow::adapters::tensorflow::GetTensorFlowType(dtype, out);
   if (!status.ok()) {
@@ -118,6 +122,7 @@ class ArrowAssignSpecImpl : public arrow::ArrayVisitor {
   VISIT_PRIMITIVE(arrow::FloatArray)
   VISIT_PRIMITIVE(arrow::DoubleArray)
   VISIT_PRIMITIVE(arrow::StringArray)
+  VISIT_PRIMITIVE(arrow::BinaryArray)
 #undef VISIT_PRIMITIVE
 
   virtual arrow::Status Visit(const arrow::ListArray& array) override {
@@ -286,6 +291,17 @@ class ArrowAssignTensorImpl : public arrow::ArrayVisitor {
     return arrow::Status::OK();
   }
 
+  virtual arrow::Status Visit(const arrow::BinaryArray& array) override {
+    auto shape = out_tensor_->shape();
+    auto output_flat = out_tensor_->flat<tstring>();
+
+    for (int64 j = 0; j < shape.num_elements(); ++j) {
+      output_flat(j) = array.GetString(i_ + j);
+    }
+
+    return arrow::Status::OK();
+  }
+
  private:
   int64 i_;
   int32 curr_array_length_;

tests/test_arrow.py

@@ -1230,6 +1230,74 @@ def test_arrow_list_feather_columns(self):
 
         os.unlink(f.name)
 
+    def test_arrow_feather_dataset_binary(self):
+        """test_arrow_feather_dataset_binary"""
+        import tensorflow_io.arrow as arrow_io
+
+        from pyarrow import feather as pa_feather
+        import pyarrow as pa
+        import numpy as np
+
+        local_path = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)), "test_image", "sample.png"
+        )
+        with open(local_path, "rb") as f:
+            data = [f.read()]
+            table = pa.Table.from_arrays([data], ["data"])
+
+        chunk_size = 1000
+        with tempfile.NamedTemporaryFile(delete=False) as f:
+            pa_feather.write_feather(table, f, chunksize=chunk_size)
+
+        dataset = arrow_io.ArrowFeatherDataset(
+            f.name,
+            columns=(0,),
+            output_types=(tf.string),
+            output_shapes=([]),
+            batch_size=32,
+        )
+        print(dataset.element_spec)
+        sample = next(iter(dataset))
+        assert sample.shape == [1]
+        assert sample.dtype == tf.string
+        assert np.all(sample == data[0])
+        print(sample.dtype, sample.shape)
+
+        os.unlink(f.name)
+
+    def test_arrow_feather_dataset_string(self):
+        """test_arrow_feather_dataset_string"""
+        import tensorflow_io.arrow as arrow_io
+
+        from pyarrow import feather as pa_feather
+        import pyarrow as pa
+        import numpy as np
+        import random
+        import string
+
+        data = ["".join(random.choice(string.printable) for _ in range(10000))]
+        table = pa.Table.from_arrays([data], schema=pa.schema([("data", pa.string())]))
+
+        chunk_size = 1000
+        with tempfile.NamedTemporaryFile(delete=False) as f:
+            pa_feather.write_feather(table, f, chunksize=chunk_size)
+
+        dataset = arrow_io.ArrowFeatherDataset(
+            f.name,
+            columns=(0,),
+            output_types=(tf.string),
+            output_shapes=([]),
+            batch_size=32,
+        )
+        print(dataset.element_spec)
+        sample = next(iter(dataset))
+        assert sample.shape == [1]
+        assert sample.dtype == tf.string
+        assert np.all(sample == data[0])
+        print(sample.dtype, sample.shape)
+
+        os.unlink(f.name)
+
 
 if __name__ == "__main__":
     test.main()