Method Compute

tensorflow/core/kernels/string_split_op.cc:375–420

  }

  void Compute(OpKernelContext* ctx) override {
    const Tensor* input_tensor;
    srand((unsigned)time(NULL));

    OP_REQUIRES_OK(ctx, ctx->input("input", &input_tensor));
    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(input_tensor->shape()),
                errors::InvalidArgument("input must be a vector, got shape: ",
                                        input_tensor->shape().DebugString()));

    const auto input_vec = input_tensor->vec<string>();
    const int64 batch_size = input_vec.dimension(0);

    const Tensor* delimiter_tensor;
    OP_REQUIRES_OK(ctx, ctx->input("delimiter", &delimiter_tensor));
    OP_REQUIRES(
        ctx, TensorShapeUtils::IsScalar(delimiter_tensor->shape()),
        errors::InvalidArgument("delimiter must be a scalar, got shape: ",
                                delimiter_tensor->shape().DebugString()));
    const auto delimiter_vec = delimiter_tensor->flat<string>();
    const string& delimiter = delimiter_vec(0);

    uint64 start = 0;
    uint64 end = 0;

    if (element_cost_ == 0 && batch_size) {
      size_t sample_id = rand() % batch_size;
      std::vector<StringPiece> temp_for_warm_up =
            skip_empty_ ?
               Split(input_vec(sample_id), delimiter, str_util::SkipEmpty())
               : Split(input_vec(sample_id), delimiter, str_util::AllowEmpty());
      start = Env::Default()->NowNanos();
      temp_for_warm_up =
            skip_empty_ ?
             Split(input_vec(sample_id), delimiter, str_util::SkipEmpty())
             : Split(input_vec(sample_id), delimiter, str_util::AllowEmpty());
      end = Env::Default()->NowNanos();
      element_cost_ = end - start;
    }
    uint64 element_cost = element_cost_;

    if (element_cost * batch_size >= parallel_limit_) {
      ParallelSplit(ctx, input_vec, batch_size, delimiter);
    } else {
      SequentialSplit(ctx, input_vec, batch_size, delimiter);
    }
  }

 private:
  bool skip_empty_;