Update JSONIOTensor with __call__ so that it is possible to support column selection

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

Author: yongtang

tensorflow_io/core/python/ops/io_tensor.py

@@ -19,6 +19,7 @@
 
 import sys
 import collections
+import uuid
 
 import tensorflow as tf
 from tensorflow_io.core.python.ops import core_ops
@@ -457,7 +458,9 @@ def __init__(self,
                internal=False):
     with tf.name_scope("AudioIOTensor") as scope:
       resource, dtypes, shapes, rate = core_ops.wav_indexable_init(
-          filename, container=scope, shared_name=filename)
+          filename,
+          container=scope,
+          shared_name="%s/%s" % (filename, uuid.uuid4().hex))
       shapes = [
           tf.TensorShape(
               [None if dim < 0 else dim for dim in e.numpy() if dim != 0]
@@ -493,27 +496,46 @@ class JSONIOTensor(IOTensor):
   #=============================================================================
   def __init__(self,
                filename,
+               columns=None,
                internal=False):
     with tf.name_scope("JSONIOTensor") as scope:
+      metadata = []
+      if columns is not None:
+        metadata.extend(["column: "+column for column in columns])
       resource, dtypes, shapes, columns = core_ops.json_indexable_init(
-          filename, container=scope, shared_name=filename)
+          filename, metadata=metadata,
+          container=scope,
+          shared_name="%s/%s" % (filename, uuid.uuid4().hex))
       shapes = [
           tf.TensorShape(
               [None if dim < 0 else dim for dim in e.numpy() if dim != 0]
           ) for e in tf.unstack(shapes)]
       dtypes = [tf.as_dtype(e.numpy()) for e in tf.unstack(dtypes)]
-      spec = [tf.TensorSpec(shape, dtype) for (
-          shape, dtype) in zip(shapes, dtypes)]
+      columns = [e.numpy() for e in tf.unstack(columns)]
+      spec = [tf.TensorSpec(shape, dtype, column) for (
+          shape, dtype, column) in zip(shapes, dtypes, columns)]
       if len(spec) == 1:
         spec = spec[0]
       else:
         spec = tuple(spec)
-      columns = [e.numpy() for e in tf.unstack(columns)]
+      self._filename = filename
       super(JSONIOTensor, self).__init__(
           spec, resource, core_ops.json_indexable_get_item,
           None,
           internal=internal)
 
+  #=============================================================================
+  # Accessors
+  #=============================================================================
+
+  def column(self, name):
+    """The `TensorSpec` of column named `name`"""
+    return next(e for e in tf.nest.flatten(self.spec) if e.name == name)
+
+  def __call__(self, column):
+    """Return a new JSONIOTensor with column named `column`"""
+    return JSONIOTensor(self._filename, columns=[column], internal=True)
+
 class KafkaIOTensor(IOIterableTensor):
   """KafkaIOTensor"""
 
@@ -541,7 +563,8 @@ def func_init(data):
         """func_init"""
         resource, _, _ = core_ops.kafka_iterable_init(
             data["subscription"], metadata=data["metadata"],
-            container=scope, shared_name=subscription)
+            container=scope,
+            shared_name="%s/%s" % (subscription, uuid.uuid4().hex))
         return resource
       func_next = core_ops.kafka_iterable_next
 

tensorflow_io/json/kernels/json_kernels.cc

@@ -234,12 +234,35 @@ class JSONIndexable : public IOIndexableInterface {
       return errors::InvalidArgument("unable to read table: ", status);
     }
 
+    std::vector<string> columns;
+    for (size_t i = 0; i < metadata.size(); i++) {
+      if (metadata[i].find_first_of("column: ") == 0) {
+        columns.emplace_back(metadata[i].substr(8));
+      }
+    }
+
+    columns_index_.clear();
+    if (columns.size() == 0) {
+      for (int i = 0; i < table_->num_columns(); i++) {
+        columns_index_.push_back(i);
+      }
+    } else {
+      std::unordered_map<string, int> columns_map;
+      for (int i = 0; i < table_->num_columns(); i++) {
+        columns_map[table_->column(i)->name()] = i;
+      }
+      for (size_t i = 0; i < columns.size(); i++) {
+        columns_index_.push_back(columns_map[columns[i]]);
+      }
+    }
+
     dtypes_.clear();
     shapes_.clear();
     columns_.clear();
-    for (int i = 0; i < table_->num_columns(); i++) {
+    for (size_t i = 0; i < columns_index_.size(); i++) {
+      int column_index = columns_index_[i];
       ::tensorflow::DataType dtype;
-      switch (table_->column(i)->type()->id()) {
+      switch (table_->column(column_index)->type()->id()) {
       case ::arrow::Type::BOOL:
         dtype = ::tensorflow::DT_BOOL;
         break;
@@ -296,7 +319,7 @@ class JSONIndexable : public IOIndexableInterface {
       }
       dtypes_.push_back(dtype);
       shapes_.push_back(TensorShape({static_cast<int64>(table_->num_rows())}));
-      columns_.push_back(table_->column(i)->name());
+      columns_.push_back(table_->column(column_index)->name());
     }
 
     return Status::OK();
@@ -329,7 +352,8 @@ class JSONIndexable : public IOIndexableInterface {
       return errors::InvalidArgument("step ", step, " is not supported");
     }
     for (size_t i = 0; i < tensors.size(); i++) {
-      std::shared_ptr<::arrow::Column> slice = table_->column(i)->Slice(start, stop);
+      int column_index = columns_index_[i];
+      std::shared_ptr<::arrow::Column> slice = table_->column(column_index)->Slice(start, stop);
 
       #define PROCESS_TYPE(TTYPE,ATYPE) { \
           int64 curr_index = 0; \
@@ -398,6 +422,7 @@ class JSONIndexable : public IOIndexableInterface {
   std::vector<DataType> dtypes_;
   std::vector<TensorShape> shapes_;
   std::vector<string> columns_;
+  std::vector<int> columns_index_;
 };
 
 REGISTER_KERNEL_BUILDER(Name("JSONIndexableInit").Device(DEVICE_CPU),

tensorflow_io/json/ops/json_ops.cc

@@ -21,6 +21,7 @@ namespace tensorflow {
 
 REGISTER_OP("JSONIndexableInit")
   .Input("input: string")
+  .Input("metadata: string")
   .Output("output: resource")
   .Output("dtypes: int64")
   .Output("shapes: int64")

tests/test_json/feature.ndjson

@@ -0,0 +1,2 @@
+{ "floatfeature": 1.1, "integerfeature": 2 }
+{ "floatfeature": 2.1, "integerfeature": 3 }

tests/test_json/label.ndjson

@@ -0,0 +1,2 @@
+{ "floatlabel": 2.2, "integerlabel": 3 }
+{ "floatlabel": 1.2, "integerlabel": 3 }

tests/test_json_eager.py

@@ -23,8 +23,65 @@
 import tensorflow as tf
 if not (hasattr(tf, "version") and tf.version.VERSION.startswith("2.")):
   tf.compat.v1.enable_eager_execution()
+import tensorflow_io as tfio  # pylint: disable=wrong-import-position
 import tensorflow_io.json as json_io  # pylint: disable=wrong-import-position
 
+def test_io_tensor_json():
+  """Test case for tfio.IOTensor.from_json."""
+  x_test = [[1.1, 2], [2.1, 3]]
+  y_test = [[2.2, 3], [1.2, 3]]
+  feature_filename = os.path.join(
+      os.path.dirname(os.path.abspath(__file__)),
+      "test_json",
+      "feature.ndjson")
+  feature_filename = "file://" + feature_filename
+  label_filename = os.path.join(
+      os.path.dirname(os.path.abspath(__file__)),
+      "test_json",
+      "label.ndjson")
+  label_filename = "file://" + label_filename
+
+  features = tfio.IOTensor.from_json(feature_filename)
+  assert features.column("floatfeature").dtype == tf.float64
+  assert features.column("integerfeature").dtype == tf.int64
+
+  labels = tfio.IOTensor.from_json(label_filename)
+  assert labels.column("floatlabel").dtype == tf.float64
+  assert labels.column("integerlabel").dtype == tf.int64
+
+  float_feature = features("floatfeature")
+  integer_feature = features("integerfeature")
+  float_label = labels("floatlabel")
+  integer_label = labels("integerlabel")
+
+  for i in range(2):
+    v_x = x_test[i]
+    v_y = y_test[i]
+    assert v_x[0] == float_feature[i].numpy()
+    assert v_x[1] == integer_feature[i].numpy()
+    assert v_y[0] == float_label[i].numpy()
+    assert v_y[1] == integer_label[i].numpy()
+
+  feature_dataset = features.to_dataset()
+
+  label_dataset = labels.to_dataset()
+
+  dataset = tf.data.Dataset.zip((
+      feature_dataset,
+      label_dataset
+  ))
+
+  i = 0
+  for (j_x, j_y) in dataset:
+    v_x = x_test[i]
+    v_y = y_test[i]
+    for index, x in enumerate(j_x):
+      assert v_x[index] == x.numpy()
+    for index, y in enumerate(j_y):
+      assert v_y[index] == y.numpy()
+    i += 1
+  assert i == len(y_test)
+
 def test_json_dataset():
   """Test case for JSON Dataset.
   """