com.linkedin.dagli.dag.AbstractDynamicDAGResult Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of core Show documentation

DAG-oriented machine learning framework for bug-resistant, readable, efficient, maintainable and trivially deployable models in Java and other JVM languages

There is a newer version: 15.0.0-beta9

Show newest version

package com.linkedin.dagli.dag;

import com.linkedin.dagli.handle.ProducerHandle;
import com.linkedin.dagli.objectio.ObjectReader;
import com.linkedin.dagli.objectio.tuple.TupleReader;
import com.linkedin.dagli.objectio.WrappedObjectReader;
import com.linkedin.dagli.producer.Producer;
import com.linkedin.dagli.util.invariant.Arguments;
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
import java.util.Arrays;


class AbstractDynamicDAGResult extends WrappedObjectReader {
  private final ObjectReader[] _outputs;
  private final Object2IntOpenHashMap> _outputToIndexMap;

  /**
   * Creates a new instance with the provided outputs
   * @param outputs the output value {@link ObjectReader}s
   * @param outputToIndexMap a map of output handles to their indices in the output list
   */
  @SuppressWarnings("unchecked")
  public AbstractDynamicDAGResult(ObjectReader[] outputs, Object2IntOpenHashMap> outputToIndexMap) {
    super((ObjectReader) (outputs.length == 1 ? outputs[0] : new TupleReader<>(outputs)));
    _outputs = outputs;
    _outputToIndexMap = outputToIndexMap;
  }

  /**
   * Gets the {@link ObjectReader} containing the results for the specified output of the dynamic DAG.
   *
   * @param output the output whose results are sought
   * @param  the type of result of the specified output
   * @return the {@link ObjectReader} containing the results for the specified output of the dynamic DAG
   */
  public  ObjectReader getResult(Producer output) {
    return getResult(output.internalAPI().getHandle());
  }

  /**
   * Gets the {@link ObjectReader} containing the results for the specified output of the dynamic DAG.
   *
   * @param outputHandle the handle of the output whose results are sought
   * @param  the type of result of the specified output
   * @return the {@link ObjectReader} containing the results for the specified output of the dynamic DAG
   */
  @SuppressWarnings("unchecked")
  public  ObjectReader getResult(ProducerHandle> outputHandle) {
    int index = _outputToIndexMap.getOrDefault(outputHandle, -1);
    Arguments.check(index >= 0, "The given output producer is not among those used to create the dynamic DAG");
    return (ObjectReader) _outputs[index];
  }

  @Override
  public void close() {
    Arrays.stream(_outputs).forEach(ObjectReader::close);
  }
}