message: [tf.data] Delete unused iterators in ChooseFastestDataset after the fastest one has been determined. This frees up resources used by the other iterators, e.g. buffers.
diff: @@ -217,8 +217,7 @@ class ChooseFastestDatasetOp : public DatasetOpKernel {
  }
  return threads[0].result->status;
  }
  - return input_impls_[fastest_index_]->GetNext(ctx, out_tensors,
  - end_of_sequence);
  + return fastest_input_impl_->GetNext(ctx, out_tensors, end_of_sequence);
  }
  protected:
  @@ -232,7 +231,14 @@ class ChooseFastestDatasetOp : public DatasetOpKernel {
  // from scratch.
  Status SaveInternal(IteratorStateWriter* writer) override {
  mutex_lock l(mu_);
  - if (input_impls_.empty()) {
  + TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("experiment_counter"),
  + experiment_counter_));
  +
  + TF_RETURN_IF_ERROR(
  + writer->WriteScalar(full_name("fastest_index"), fastest_index_));
  + if (fastest_index_ != -1) {
  + TF_RETURN_IF_ERROR(SaveInput(writer, fastest_input_impl_));
  + } else if (input_impls_.empty()) {
  TF_RETURN_IF_ERROR(
  writer->WriteScalar(full_name("input_impls_empty"), ""));
  } else {
  @@ -240,17 +246,22 @@ class ChooseFastestDatasetOp : public DatasetOpKernel {
  TF_RETURN_IF_ERROR(SaveInput(writer, input_impl));
  }
  }
  - TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("experiment_counter"),
  - experiment_counter_));
  - TF_RETURN_IF_ERROR(
  - writer->WriteScalar(full_name("fastest_index"), fastest_index_));
  return Status::OK();
  }
  Status RestoreInternal(IteratorContext* ctx,
  IteratorStateReader* reader) override {
  mutex_lock l(mu_);
  - if (reader->Contains(full_name("input_impls_empty"))) {
  + TF_RETURN_IF_ERROR(reader->ReadScalar(full_name("experiment_counter"),
  + &experiment_counter_));
  + TF_RETURN_IF_ERROR(
  + reader->ReadScalar(full_name("fastest_index"), &fastest_index_));
  + if (fastest_index_ != -1) {
  + TF_RETURN_IF_ERROR(dataset()->inputs_[fastest_index_]->MakeIterator(
  + ctx, strings::StrCat(prefix(), "_", fastest_index_),
  + &fastest_input_impl_));
  + TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, fastest_input_impl_));
  + } else if (reader->Contains(full_name("input_impls_empty"))) {
  input_impls_.clear();
  } else {
  DCHECK_EQ(input_impls_.size(), dataset()->inputs_.size());
  @@ -258,10 +269,6 @@ class ChooseFastestDatasetOp : public DatasetOpKernel {
  TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl));
  }
  }
  - TF_RETURN_IF_ERROR(reader->ReadScalar(full_name("experiment_counter"),
  - &experiment_counter_));
  - TF_RETURN_IF_ERROR(
  - reader->ReadScalar(full_name("fastest_index"), &fastest_index_));
  return Status::OK();
  }
  @@ -279,6 +286,7 @@ class ChooseFastestDatasetOp : public DatasetOpKernel {
  };
  std::vector<std::unique_ptr<IteratorBase>> input_impls_;
  + std::unique_ptr<IteratorBase> fastest_input_impl_;
  // For tracking the time taken for each input's iterations.
  std::vector<histogram::Histogram> histograms_;
  @@ -326,6 +334,9 @@ class ChooseFastestDatasetOp : public DatasetOpKernel {
  fastest_index_ = i;
  }
  }
  +
  + fastest_input_impl_ = std::move(input_impls_[fastest_index_]);
  + input_impls_.clear(); // Delete the unused iterators.
  }
  }; // class Iterator
  