Method SetIteratorFromDataset

tensorflow/core/kernels/data/iterator_ops.cc:138–174

}

Status IteratorResource::SetIteratorFromDataset(OpKernelContext* ctx,
                                                DatasetBase* dataset) {
  std::shared_ptr<State> new_state;
  {
    tf_shared_lock l(mu_);
    new_state = std::make_shared<State>(
        iterator_state_->flib_def, iterator_state_->pflr, iterator_state_->flr,
        /*iterator=*/nullptr);
  }
  // Create new iterator.
  std::unique_ptr<IteratorBase> iterator;
  IteratorContext::Params params(ctx);
  params.flr = new_state->flr;
  params.function_handle_cache = new_state->function_handle_cache.get();
  params.resource_mgr = &new_state->resource_mgr;
  params.thread_factory = unbounded_thread_pool_.get_thread_factory();
  params.thread_pool = &unbounded_thread_pool_;
  params.cancellation_manager = &new_state->cancellation_manager;
  std::function<void()> deregister_fn;
  TF_RETURN_IF_ERROR(ConnectCancellationManagers(ctx->cancellation_manager(),
                                                 params.cancellation_manager,
                                                 &deregister_fn));
  {
    auto cleanup = gtl::MakeCleanup(std::move(deregister_fn));
    TF_RETURN_IF_ERROR(dataset->MakeIterator(IteratorContext(std::move(params)),
                                             "Iterator", &iterator));
    TF_RETURN_IF_ERROR(
        VerifyTypesMatch(output_dtypes_, iterator->output_dtypes()));
    TF_RETURN_IF_ERROR(
        VerifyShapesCompatible(output_shapes_, iterator->output_shapes()));
    std::swap(new_state->iterator, iterator);
  }

  mutex_lock l(mu_);
  std::swap(iterator_state_, new_state);
  return Status::OK();
}

namespace {