Method ComputeAsync

tf_euler/kernels/get_edge_binary_feature_op.cc:61–122

};

void GetEdgeBinaryFeature::ComputeAsync(OpKernelContext* ctx,
                                        DoneCallback done) {
  auto edges = ctx->input(0);
  auto& shape = edges.shape();
  OP_REQUIRES(ctx, shape.dim_size(1) == 3,
              errors::InvalidArgument(
                  "Input `edges` must be a matrix with shape [n, 3]"));

  std::vector<Tensor*> outputs(N_, nullptr);
  TensorShape output_shape;
  output_shape.AddDim(shape.dim_size(0));
  for (auto i = 0; i < N_; ++i) {
    OP_REQUIRES_OK(ctx, ctx->allocate_output(i, output_shape, &outputs[i]));
  }

  auto edges_flat = edges.flat<int64>();
  size_t edges_size = edges_flat.size();
  size_t edge_num  = edges_size / 3;

  auto query = new euler::Query(query_str_);
  auto t_edges = query->AllocInput("edges", {edge_num, 3}, euler::kUInt64);
  std::copy(edges_flat.data(),
            edges_flat.data() + edges_size, t_edges->Raw<int64_t>());
  for (size_t i = 0; i < feature_names_.size(); ++i) {
    auto t_fid = query->AllocInput("__" + feature_names_[i], {1},
                                   euler::kString);
    *(t_fid->Raw<std::string*>()[0]) = "binary_" + feature_names_[i];
  }

  auto callback = [outputs, done, query, edge_num, this]() {
    std::stringstream ss;
    auto results_map = query->GetResult(res_names_);
    for (size_t i = 0 ; i < feature_names_.size(); ++i) {
      ss.str("");
      ss << "fea:" << i * 2;
      std::string fea_idx = ss.str();

      ss.str("");
      ss << "fea:" << i * 2 + 1;
      std::string fea_val = ss.str();

      if (results_map[fea_idx]->NumElements() != edge_num * 2) {
        EULER_LOG(FATAL) << "Binary Feature Result Index Num Error:" <<
            results_map[fea_idx]->NumElements() << "Expect: " << edge_num * 2;
      }
      for (size_t j = 0; j < edge_num; ++j) {
        size_t start = results_map[fea_idx]->Raw<int32_t>()[j * 2];
        size_t end = results_map[fea_idx]->Raw<int32_t>()[j * 2 + 1];
        auto data = outputs[i]->flat<tensorflow::string>();
        std::string f_v(end - start, 0);
        std::copy(results_map[fea_val]->Raw<char>() + start,
                  results_map[fea_val]->Raw<char>() + end,
                  f_v.begin());
        data(j) = f_v;
      }
    }
    delete query;
    done();