Add from_json, and allows multiple columns

Signed-off-by: Yong Tang <[email protected]>

Author: yongtang

WORKSPACE

@@ -542,3 +542,13 @@ http_archive(
         "https://github.com/google/double-conversion/archive/3992066a95b823efc8ccc1baf82a1cfc73f6e9b8.zip",
     ],
 )
+
+http_archive(
+    name = "rapidjson",
+    build_file = "//third_party:rapidjson.BUILD",
+    sha256 = "bf7ced29704a1e696fbccf2a2b4ea068e7774fa37f6d7dd4039d0787f8bed98e",
+    strip_prefix = "rapidjson-1.1.0",
+    urls = [
+        "https://github.com/miloyip/rapidjson/archive/v1.1.0.tar.gz",
+    ],
+)

tensorflow_io/arrow/kernels/arrow_kernels.h

@@ -30,7 +30,8 @@ class ArrowRandomAccessFile : public ::arrow::io::RandomAccessFile {
 public:
   explicit ArrowRandomAccessFile(tensorflow::RandomAccessFile *file, int64 size)
     : file_(file)
-    , size_(size) { }
+    , size_(size)
+    , position_(0) { }
 
   ~ArrowRandomAccessFile() {}
   arrow::Status Close() override {
@@ -40,13 +41,21 @@ class ArrowRandomAccessFile : public ::arrow::io::RandomAccessFile {
     return false;
   }
   arrow::Status Tell(int64_t* position) const override {
-    return arrow::Status::NotImplemented("Tell");
+    *position = position_;
+    return arrow::Status::OK();
   }
   arrow::Status Seek(int64_t position) override {
     return arrow::Status::NotImplemented("Seek");
   }
   arrow::Status Read(int64_t nbytes, int64_t* bytes_read, void* out) override {
-    return arrow::Status::NotImplemented("Read (void*)");
+    StringPiece result;
+    Status status = file_->Read(position_, nbytes, &result, (char*)out);
+    if (!(status.ok() || errors::IsOutOfRange(status))) {
+        return arrow::Status::IOError(status.error_message());
+    }
+    *bytes_read = result.size();
+    position_ += (*bytes_read);
+    return arrow::Status::OK();
   }
   arrow::Status Read(int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) override {
     return arrow::Status::NotImplemented("Read (Buffer*)");
@@ -81,6 +90,7 @@ class ArrowRandomAccessFile : public ::arrow::io::RandomAccessFile {
 private:
   tensorflow::RandomAccessFile* file_;
   int64 size_;
+  int64 position_;
 };
 }  // namespace data
 }  // namespace tensorflow

tensorflow_io/audio/ops/audio_ops.cc

@@ -42,9 +42,19 @@ REGISTER_OP("WAVIndexableGetItem")
   .Input("stop: int64")
   .Input("step: int64")
   .Output("output: dtype")
-  .Attr("dtype: {int8, int16, int32}")
+  .Attr("dtype: list(type) >= 1")
+  .Attr("shape: list(shape) >= 1")
   .SetShapeFn([](shape_inference::InferenceContext* c) {
-    c->set_output(0, c->UnknownShape());
+    std::vector<PartialTensorShape> shape;
+    TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape));
+    if (shape.size() != c->num_outputs()) {
+      return errors::InvalidArgument("`shape` must be the same length as `types` (", shape.size(), " vs. ", c->num_outputs());
+    }
+    for (size_t i = 0; i < shape.size(); ++i) {
+      shape_inference::ShapeHandle entry;
+      TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(shape[i], &entry));
+      c->set_output(static_cast<int64>(i), entry);
+    }
     return Status::OK();
    });
 

tensorflow_io/audio/python/ops/audio_ops.py

@@ -21,10 +21,10 @@
 from tensorflow_io.core.python.ops import data_ops
 from tensorflow_io.core.python.ops import io_tensor
 
-class WAVDataset(data_ops.BaseDataset):
+class WAVDataset(tf.compat.v2.data.Dataset):
   """A WAV Dataset"""
 
-  def __init__(self, filename, batch=None, **kwargs):
+  def __init__(self, filename, batch=None):
     """Create a WAVDataset.
 
     Args:
@@ -35,28 +35,16 @@ def __init__(self, filename, batch=None, **kwargs):
       raise NotImplementedError("WAVDataset only support eager mode")
 
     self._wav = io_tensor.IOTensor.from_audio(filename)
+    self._dataset = self._wav.to_dataset()
+    super(WAVDataset, self).__init__(
+        self._dataset._variant_tensor) # pylint: disable=protected-access
 
-    dtype = self._wav.dtype
-    shape = self._wav.shape[1:]
-    start = 0
-    stop = self._wav.shape[0]
-
-    # capacity is the rough count for each chunk in dataset
-    capacity = kwargs.get("capacity", 65536)
-    entry_start = list(range(start, stop, capacity))
-    entry_stop = entry_start[1:] + [stop]
-    dataset = data_ops.BaseDataset.from_tensor_slices(
-        (tf.constant(entry_start, tf.int64), tf.constant(entry_stop, tf.int64))
-    ).map(lambda start, stop: self._wav.__getitem__(slice(start, stop)))
-
-    dataset = dataset.apply(tf.data.experimental.unbatch())
-    if batch != 0:
-      dataset = dataset.batch(batch)
-      shape = tf.TensorShape([None]).concatenate(shape)
+  def _inputs(self):
+    return []
 
-    self._dataset = dataset
-    super(WAVDataset, self).__init__(
-        self._dataset._variant_tensor, [dtype], [shape]) # pylint: disable=protected-access
+  @property
+  def _element_structure(self):
+    return self._dataset._element_structure # pylint: disable=protected-access
 
 class AudioDataset(data_ops.Dataset):
   """A Audio File Dataset that reads the audio file."""