Status FillTensorMapFromInstancesList()

in tensorflow_serving/util/json_tensor.cc [462:550]
73 lines of code
15 McCabe index (conditional complexity)

Status FillTensorMapFromInstancesList(
    const rapidjson::Value::MemberIterator& itr,
    const ::google::protobuf::Map<string, tensorflow::TensorInfo>& tensorinfo_map,
    ::google::protobuf::Map<string, TensorProto>* tensor_map) {
  // "instances" array can either be a plain list or list of objects (for named
  // tensors) but not a mix of both. Each object must have one key for each
  // named tensor.
  if (!itr->value[0].IsObject() && tensorinfo_map.size() > 1) {
    return errors::InvalidArgument(
        "instances is a plain list, but expecting list of objects as multiple "
        "input tensors required as per tensorinfo_map");
  }

  auto IsElementObject = [](const rapidjson::Value& val) {
    return val.IsObject() && !IsValBase64Object(val);
  };

  const bool elements_are_objects = IsElementObject(itr->value[0]);

  std::set<string> input_names;
  for (const auto& kv : tensorinfo_map) input_names.insert(kv.first);

  // Add each element of "instances" array to tensor.
  //
  // Each element must yield one tensor of same shape and size. All elements get
  // batched into one tensor with the first dimension equal to the number of
  // elements in the instances array.
  tensor_map->clear();
  ::google::protobuf::Map<string, int> size_map;
  ::google::protobuf::Map<string, TensorShapeProto> shape_map;
  int tensor_count = 0;
  for (const auto& elem : itr->value.GetArray()) {
    if (elements_are_objects) {
      if (!IsElementObject(elem)) {
        return errors::InvalidArgument("Expecting object but got list at item ",
                                       tensor_count, " of input list");
      }
      std::set<string> object_keys;
      for (const auto& kv : elem.GetObject()) {
        const string& name = kv.name.GetString();
        object_keys.insert(name);
        const auto status = AddInstanceItem(kv.value, name, tensorinfo_map,
                                            &size_map, &shape_map, tensor_map);
        if (!status.ok()) {
          return errors::InvalidArgument(
              "Failed to process element: ", tensor_count, " key: ", name,
              " of 'instances' list. Error: ", status.ToString());
        }
      }
      if (input_names != object_keys) {
        return errors::InvalidArgument(
            "Failed to process element: ", tensor_count,
            " of 'instances' list. JSON object: ", JsonValueToString(elem),
            " keys must be equal to: ", absl::StrJoin(input_names, ","));
      }
    } else {
      if (IsElementObject(elem)) {
        return errors::InvalidArgument(
            "Expecting value/list but got object at item ", tensor_count,
            " of input list");
      }
      const auto& name = tensorinfo_map.begin()->first;
      const auto status = AddInstanceItem(elem, name, tensorinfo_map, &size_map,
                                          &shape_map, tensor_map);
      if (!status.ok()) {
        return errors::InvalidArgument(
            "Failed to process element: ", tensor_count,
            " of 'instances' list. Error: ", status.ToString());
      }
    }
    tensor_count++;
  }

  // Now that all individual tensors from "instances" array are added,
  // fix the resulting shape of the final tensor.
  for (auto& kv : *tensor_map) {
    const string& name = kv.first;
    auto* tensor = &kv.second;
    tensor->set_dtype(tensorinfo_map.at(name).dtype());
    const auto& shape = shape_map.at(name);
    auto* output_shape = tensor->mutable_tensor_shape();
    output_shape->Clear();
    output_shape->add_dim()->set_size(tensor_count);
    for (const auto& d : shape.dim())
      output_shape->add_dim()->set_size(d.size());
  }

  return Status::OK();
}