Add decode_json support

tensorflow_io/core/BUILD

@@ -368,6 +368,21 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "serialization_ops",
+    srcs = [
+        "kernels/serialization_kernels.cc",
+        "ops/serialization_ops.cc",
+    ],
+    copts = tf_io_copts(),
+    linkstatic = True,
+    deps = [
+        "//tensorflow_io/core:dataset_ops",
+        "@rapidjson",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "grpc_ops",
     srcs = [
@@ -424,6 +439,7 @@ cc_binary(
         "//tensorflow_io/core:parquet_ops",
         "//tensorflow_io/core:pcap_ops",
         "//tensorflow_io/core:pubsub_ops",
+        "//tensorflow_io/core:serialization_ops",
         "//tensorflow_io/core:text_ops",
         "//tensorflow_io/core/azure:azfs_ops",
         "//tensorflow_io/core/oss:oss_ops",

tensorflow_io/core/kernels/serialization_kernels.cc

@@ -0,0 +1,136 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/resource_op_kernel.h"
+
+#include "rapidjson/document.h"
+#include "rapidjson/pointer.h"
+
+namespace tensorflow {
+namespace data {
+namespace {
+
+class DecodeJSONOp : public OpKernel {
+ public:
+  explicit DecodeJSONOp(OpKernelConstruction* context) : OpKernel(context) {
+    env_ = context->env();
+    OP_REQUIRES_OK(context, context->GetAttr("shapes", &shapes_));
+  }
+
+  void Compute(OpKernelContext* context) override {
+    // TODO: support batch (1-D) input
+    const Tensor* input_tensor;
+    OP_REQUIRES_OK(context, context->input("input", &input_tensor));
+    const string& input = input_tensor->scalar<string>()();
+
+    const Tensor* names_tensor;
+    OP_REQUIRES_OK(context, context->input("names", &names_tensor));
+
+    OP_REQUIRES(context, (names_tensor->NumElements() == shapes_.size()),
+                errors::InvalidArgument(
+                    "shapes and names should have same number: ",
+                    shapes_.size(), " vs. ", names_tensor->NumElements()));
+    rapidjson::Document d;
+    d.Parse(input.c_str());
+    OP_REQUIRES(context, d.IsObject(),
+                errors::InvalidArgument("not a valid JSON object"));
+    for (size_t i = 0; i < shapes_.size(); i++) {
+      Tensor* value_tensor;
+      OP_REQUIRES_OK(context,
+                     context->allocate_output(i, shapes_[i], &value_tensor));
+      rapidjson::Value* entry =
+          rapidjson::Pointer(names_tensor->flat<string>()(i).c_str()).Get(d);
+      OP_REQUIRES(context, (entry != nullptr),
+                  errors::InvalidArgument("no value for ",
+                                          names_tensor->flat<string>()(i)));
+      if (entry->IsArray()) {
+        OP_REQUIRES(context, entry->Size() == value_tensor->NumElements(),
+                    errors::InvalidArgument(
+                        "number of elements in JSON does not match spec: ",
+                        entry->Size(), " vs. ", value_tensor->NumElements()));
+
+        switch (value_tensor->dtype()) {
+          case DT_INT32:
+            for (int64 j = 0; j < entry->Size(); j++) {
+              value_tensor->flat<int32>()(j) = (*entry)[j].GetInt();
+            }
+            break;
+          case DT_INT64:
+            for (int64 j = 0; j < entry->Size(); j++) {
+              value_tensor->flat<int64>()(j) = (*entry)[j].GetInt64();
+            }
+            break;
+          case DT_FLOAT:
+            for (int64 j = 0; j < entry->Size(); j++) {
+              value_tensor->flat<float>()(j) = (*entry)[j].GetDouble();
+            }
+            break;
+          case DT_DOUBLE:
+            for (int64 j = 0; j < entry->Size(); j++) {
+              value_tensor->flat<double>()(j) = (*entry)[j].GetDouble();
+            }
+            break;
+          case DT_STRING:
+            for (int64 j = 0; j < entry->Size(); j++) {
+              value_tensor->flat<string>()(j) = (*entry)[j].GetString();
+            }
+            break;
+          default:
+            OP_REQUIRES(
+                context, false,
+                errors::InvalidArgument("data type not supported: ",
+                                        DataTypeString(value_tensor->dtype())));
+            break;
+        }
+
+      } else {
+        switch (value_tensor->dtype()) {
+          case DT_INT32:
+            value_tensor->scalar<int32>()() = entry->GetInt();
+            break;
+          case DT_INT64:
+            value_tensor->scalar<int64>()() = entry->GetInt64();
+            break;
+          case DT_FLOAT:
+            value_tensor->scalar<float>()() = entry->GetDouble();
+            break;
+          case DT_DOUBLE:
+            value_tensor->scalar<double>()() = entry->GetDouble();
+            break;
+          case DT_STRING:
+            value_tensor->scalar<string>()() = entry->GetString();
+            break;
+          default:
+            OP_REQUIRES(
+                context, false,
+                errors::InvalidArgument("data type not supported: ",
+                                        DataTypeString(value_tensor->dtype())));
+            break;
+        }
+      }
+    }
+  }
+
+ private:
+  mutable mutex mu_;
+  Env* env_ GUARDED_BY(mu_);
+  std::vector<TensorShape> shapes_ GUARDED_BY(mu_);
+};
+REGISTER_KERNEL_BUILDER(Name("IO>DecodeJSON").Device(DEVICE_CPU), DecodeJSONOp);
+
+}  // namespace
+}  // namespace data
+}  // namespace tensorflow

tensorflow_io/core/ops/serialization_ops.cc

@@ -0,0 +1,52 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+namespace tensorflow {
+namespace io {
+namespace {
+
+REGISTER_OP("IO>DecodeJSON")
+    .Input("input: string")
+    .Input("names: string")
+    .Output("value: dtypes")
+    .Attr("shapes: list(shape)")
+    .Attr("dtypes: list(type)")
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      // TODO: support batch (1-D) input
+      shape_inference::ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 0, &unused));
+      std::vector<TensorShape> shapes;
+      TF_RETURN_IF_ERROR(c->GetAttr("shapes", &shapes));
+      if (shapes.size() != c->num_outputs()) {
+        return errors::InvalidArgument(
+            "shapes and types should be the same: ", shapes.size(), " vs. ",
+            c->num_outputs());
+      }
+      for (size_t i = 0; i < shapes.size(); ++i) {
+        shape_inference::ShapeHandle shape;
+        TF_RETURN_IF_ERROR(
+            c->MakeShapeFromPartialTensorShape(shapes[i], &shape));
+        c->set_output(static_cast<int64>(i), shape);
+      }
+      return Status::OK();
+    });
+
+}  // namespace
+}  // namespace io
+}  // namespace tensorflow

tensorflow_io/core/python/api/experimental/__init__.py

@@ -18,6 +18,7 @@
 from tensorflow_io.core.python.experimental.io_tensor import IOTensor
 from tensorflow_io.core.python.experimental.io_layer import IOLayer
 
+from tensorflow_io.core.python.api.experimental import serialization
 from tensorflow_io.core.python.api.experimental import ffmpeg
 from tensorflow_io.core.python.api.experimental import image
 from tensorflow_io.core.python.api.experimental import text

tensorflow_io/core/python/api/experimental/serialization.py

@@ -0,0 +1,17 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""tensorflow_io.experimental.serialization"""
+
+from tensorflow_io.core.python.experimental.serialization_ops import decode_json # pylint: disable=unused-import

tensorflow_io/core/python/experimental/serialization_ops.py

@@ -0,0 +1,73 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Serialization Ops."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+from tensorflow_io.core.python.ops import core_ops
+
+# _NamedTensorSpec allows adding a `named` key while traversing,
+# so that it is possible to build up the `/R/Foo` JSON Pointers.
+class _NamedTensorSpec(tf.TensorSpec):
+  """_NamedTensorSpec"""
+  def named(self, named=None):
+    if named is not None:
+      self._named = named
+    return self._named
+
+# named_spec updates named field for JSON Pointers while traversing.
+def named_spec(specs, name=''):
+  """named_spec"""
+  if isinstance(specs, _NamedTensorSpec):
+    specs.named(name)
+    return
+
+  if isinstance(specs, dict):
+    for k in specs.keys():
+      named_spec(specs[k], "{}/{}".format(name, k))
+    return
+
+  for k, _ in enumerate(specs):
+    named_spec(specs[k], "{}/{}".format(name, k))
+  return
+
+
+def decode_json(json, specs, name=None):
+  """
+  Decode JSON string into Tensors.
+
+  TODO: support batch (1-D) input
+
+  Args:
+    json: A String Tensor. The JSON strings to decode.
+    specs: A structured TensorSpecs describing the signature
+      of the JSON elements.
+    name: A name for the operation (optional).
+
+  Returns:
+    A structured Tensors.
+  """
+  # Make a copy of specs to keep the original specs
+  named = tf.nest.map_structure(lambda e: _NamedTensorSpec(e.shape, e.dtype), specs)
+  named_spec(named)
+  named = tf.nest.flatten(named)
+  names = [e.named() for e in named]
+  shapes = [e.shape for e in named]
+  dtypes = [e.dtype for e in named]
+
+  values = core_ops.io_decode_json(json, names, shapes, dtypes, name=name)
+  return tf.nest.pack_sequence_as(specs, values)