Import GetTensorFlowType and GetArrowType

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

Author: yongtang

tensorflow_io/arrow/kernels/arrow_dataset_ops.cc

@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 
 #include "arrow/api.h"
-#include "arrow/adapters/tensorflow/convert.h"
 #include "arrow/ipc/api.h"
 #include "arrow/util/io-util.h"
 #include "tensorflow/core/framework/dataset.h"
@@ -99,8 +98,10 @@ class ArrowColumnTypeChecker : public arrow::TypeVisitor {
   // Check scalar types with arrow::adapters::tensorflow
   arrow::Status CheckScalarType(std::shared_ptr<arrow::DataType> scalar_type) {
     DataType converted_type;
-    ARROW_RETURN_NOT_OK(arrow::adapters::tensorflow::GetTensorFlowType(
-        scalar_type, &converted_type));
+    ::tensorflow::Status status = GetTensorFlowType(scalar_type, &converted_type);
+    if (!status.ok()) {
+      return ::arrow::Status::Invalid(status);
+    }
     if (converted_type != expected_type_) {
       return arrow::Status::TypeError(
           "Arrow type mismatch: expected dtype=" +
@@ -523,11 +524,7 @@ class ArrowOpKernelBase : public DatasetOpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("output_shapes", &output_shapes_));
     for (const DataType& dt : output_types_) {
       std::shared_ptr<arrow::DataType> arrow_type;
-      auto status = arrow::adapters::tensorflow::GetArrowType(dt, &arrow_type);
-      OP_REQUIRES(ctx, status.ok(),
-                  errors::InvalidArgument(
-                      "Arrow type is unsupported for output_type dtype=" +
-                      std::to_string(dt)));
+      OP_REQUIRES_OK(ctx, GetArrowType(dt, &arrow_type));
     }
     for (const PartialTensorShape& pts : output_shapes_) {
       OP_REQUIRES(ctx, -1 <= pts.dims() && pts.dims() <= 2,

tensorflow_io/arrow/kernels/arrow_kernels.cc

@@ -19,9 +19,27 @@ limitations under the License.
 #include "arrow/ipc/feather.h"
 #include "arrow/ipc/feather_generated.h"
 #include "arrow/buffer.h"
+#include "arrow/adapters/tensorflow/convert.h"
 
 namespace tensorflow {
 namespace data {
+
+Status GetTensorFlowType(std::shared_ptr<::arrow::DataType> dtype, ::tensorflow::DataType* out) {
+  ::arrow::Status status = ::arrow::adapters::tensorflow::GetTensorFlowType(dtype, out);
+  if (!status.ok()) {
+    return errors::InvalidArgument("arrow data type ", dtype, " is not supported: ", status);
+  }
+  return Status::OK();
+}
+
+Status GetArrowType(::tensorflow::DataType dtype, std::shared_ptr<::arrow::DataType>* out) {
+  ::arrow::Status status = ::arrow::adapters::tensorflow::GetArrowType(dtype, out);
+  if (!status.ok()) {
+    return errors::InvalidArgument("tensorflow data type ", dtype, " is not supported: ", status);
+  }
+  return Status::OK();
+}
+
 namespace {
 
 class ListFeatherColumnsOp : public OpKernel {

tensorflow_io/arrow/kernels/arrow_kernels.h

@@ -19,10 +19,14 @@ limitations under the License.
 #include "kernels/stream.h"
 #include "arrow/io/api.h"
 #include "arrow/buffer.h"
+#include "arrow/type.h"
 
 namespace tensorflow {
 namespace data {
 
+Status GetTensorFlowType(std::shared_ptr<::arrow::DataType> dtype, ::tensorflow::DataType* out);
+Status GetArrowType(::tensorflow::DataType dtype, std::shared_ptr<::arrow::DataType>* out);
+
 // NOTE: Both SizedRandomAccessFile and ArrowRandomAccessFile overlap
 // with another PR. Will remove duplicate once PR merged
 
@@ -92,6 +96,8 @@ class ArrowRandomAccessFile : public ::arrow::io::RandomAccessFile {
   int64 size_;
   int64 position_;
 };
+
+
 }  // namespace data
 }  // namespace tensorflow
 

tensorflow_io/json/kernels/json_kernels.cc

@@ -262,61 +262,7 @@ class JSONIndexable : public IOIndexableInterface {
     for (size_t i = 0; i < columns_index_.size(); i++) {
       int column_index = columns_index_[i];
       ::tensorflow::DataType dtype;
-      switch (table_->column(column_index)->type()->id()) {
-      case ::arrow::Type::BOOL:
-        dtype = ::tensorflow::DT_BOOL;
-        break;
-      case ::arrow::Type::UINT8:
-        dtype= ::tensorflow::DT_UINT8;
-        break;
-      case ::arrow::Type::INT8:
-        dtype= ::tensorflow::DT_INT8;
-        break;
-      case ::arrow::Type::UINT16:
-        dtype= ::tensorflow::DT_UINT16;
-        break;
-      case ::arrow::Type::INT16:
-        dtype= ::tensorflow::DT_INT16;
-        break;
-      case ::arrow::Type::UINT32:
-        dtype= ::tensorflow::DT_UINT32;
-        break;
-      case ::arrow::Type::INT32:
-        dtype= ::tensorflow::DT_INT32;
-        break;
-      case ::arrow::Type::UINT64:
-        dtype= ::tensorflow::DT_UINT64;
-        break;
-      case ::arrow::Type::INT64:
-        dtype= ::tensorflow::DT_INT64;
-        break;
-      case ::arrow::Type::HALF_FLOAT:
-        dtype= ::tensorflow::DT_HALF;
-        break;
-      case ::arrow::Type::FLOAT:
-        dtype= ::tensorflow::DT_FLOAT;
-        break;
-      case ::arrow::Type::DOUBLE:
-        dtype= ::tensorflow::DT_DOUBLE;
-        break;
-      case ::arrow::Type::STRING:
-      case ::arrow::Type::BINARY:
-      case ::arrow::Type::FIXED_SIZE_BINARY:
-      case ::arrow::Type::DATE32:
-      case ::arrow::Type::DATE64:
-      case ::arrow::Type::TIMESTAMP:
-      case ::arrow::Type::TIME32:
-      case ::arrow::Type::TIME64:
-      case ::arrow::Type::INTERVAL:
-      case ::arrow::Type::DECIMAL:
-      case ::arrow::Type::LIST:
-      case ::arrow::Type::STRUCT:
-      case ::arrow::Type::UNION:
-      case ::arrow::Type::DICTIONARY:
-      case ::arrow::Type::MAP:
-      default:
-        return errors::InvalidArgument("arrow data type is not supported: ", table_->column(i)->type()->ToString());
-      }
+      TF_RETURN_IF_ERROR(GetTensorFlowType(table_->column(column_index)->type(), &dtype));
       dtypes_.push_back(dtype);
       shapes_.push_back(TensorShape({static_cast<int64>(table_->num_rows())}));
       columns_.push_back(table_->column(column_index)->name());
@@ -347,7 +293,6 @@ class JSONIndexable : public IOIndexableInterface {
   }
 
   Status GetItem(const int64 start, const int64 stop, const int64 step, std::vector<Tensor>& tensors) override {
-    Tensor& output_tensor = tensors[0];
     if (step != 1) {
       return errors::InvalidArgument("step ", step, " is not supported");
     }