void Compute()

in tensorflow_ops/timestamp_ops_kernel.cc [301:361]


  void Compute(OpKernelContext* context) override {
    // Grab the timestamp tensor
    const Tensor& timestamp_tensor = context->input(0);
    auto timestamp = timestamp_tensor.flat<tstring>();
    // Grab the interval tensor
    const Tensor& diff_tensor = context->input(1);
    auto interval_int = diff_tensor.flat<int64_t>();
    OP_REQUIRES(
        context, interval_int.size() == timestamp.size(),
        InvalidArgument(absl::Substitute(
            "Error in $0: timestamp and interval must have the same shape, "
            "but are $1, $2",
            name(), timestamp.size(), interval_int.size())));
    // Grab the part tensor
    const Tensor& part_tensor = context->input(2);
    std::string part = part_tensor.flat<tstring>()(0);
    functions::DateTimestampPart part_enum;
    static auto* supported_parts =
        new absl::flat_hash_set<functions::DateTimestampPart>(
            {functions::MICROSECOND, functions::MILLISECOND, functions::SECOND,
             functions::MINUTE, functions::HOUR, functions::DAY});
    OP_REQUIRES_OK(context, ParseInputDateTimestampPart(
                                part, name(), &part_enum, *supported_parts));

    // Create an output tensor with the shape of the timestamp tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(
                                0, timestamp_tensor.shape(), &output_tensor));
    auto output_flat = output_tensor->flat<tstring>();

    const int N = timestamp.size();
    for (int i = 0; i < N; i++) {
      // Default time zone.
      absl::TimeZone tz = absl::UTCTimeZone();

      // Parse the timestamp.
      int64_t input_ts;
      OP_REQUIRES_OK(context,
                     ParseInputTimestamp(timestamp(i), tz, name(), &input_ts));

      absl::StatusOr<IntervalValue> interval =
          GetIntervalValue(interval_int(i), part_enum);
      OP_REQUIRES(
          context, interval.ok(),
          Internal("Error in getting interval of TimestampAdd with status: ",
                   interval.status()));
      absl::Time base_time;
      OP_REQUIRES_OK(context,
                     ToTslStatus(name(), functions::AddTimestamp(
                                             absl::FromUnixMicros(input_ts), tz,
                                             *interval, &base_time)));
      int64_t ts = absl::ToUnixMicros(base_time);

      std::string out;
      OP_REQUIRES_OK(context, FormatOutputTimestamp(ts, name(), &out));

      // Set the output value.
      output_flat(i).reserve(out.size());
      output_flat(i) = std::move(out);
    }
  }