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);
}
}