io.substrait.extension.AbstractExtensionLookup Maven / Gradle / Ivy

Go to download

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

Create a well-defined, cross-language specification for data compute operations

There is a newer version: 0.46.1

package io.substrait.extension;

import java.util.Map;

public abstract class AbstractExtensionLookup implements ExtensionLookup {
  protected final Map functionAnchorMap;
  protected final Map typeAnchorMap;

  public AbstractExtensionLookup(
      Map functionAnchorMap,
      Map typeAnchorMap) {
    this.functionAnchorMap = functionAnchorMap;
    this.typeAnchorMap = typeAnchorMap;
  }

  public SimpleExtension.ScalarFunctionVariant getScalarFunction(
      int reference, SimpleExtension.ExtensionCollection extensions) {
    var anchor = functionAnchorMap.get(reference);
    if (anchor == null) {
      throw new IllegalArgumentException(
          "Unknown function id. Make sure that the function id provided was shared in the extensions section of the plan.");
    }

    return extensions.getScalarFunction(anchor);
  }

  public SimpleExtension.WindowFunctionVariant getWindowFunction(
      int reference, SimpleExtension.ExtensionCollection extensions) {
    var anchor = functionAnchorMap.get(reference);
    if (anchor == null) {
      throw new IllegalArgumentException(
          "Unknown function id. Make sure that the function id provided was shared in the extensions section of the plan.");
    }

    return extensions.getWindowFunction(anchor);
  }

  public SimpleExtension.AggregateFunctionVariant getAggregateFunction(
      int reference, SimpleExtension.ExtensionCollection extensions) {
    var anchor = functionAnchorMap.get(reference);
    if (anchor == null) {
      throw new IllegalArgumentException(
          "Unknown function id. Make sure that the function id provided was shared in the extensions section of the plan.");
    }

    return extensions.getAggregateFunction(anchor);
  }

  public SimpleExtension.Type getType(
      int reference, SimpleExtension.ExtensionCollection extensions) {
    var anchor = typeAnchorMap.get(reference);
    if (anchor == null) {
      throw new IllegalArgumentException(
          "Unknown type id. Make sure that the type id provided was shared in the extensions section of the plan.");
    }

    return extensions.getType(anchor);
  }
}