io.substrait.relation.Aggregate 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.relation;

import io.substrait.expression.AggregateFunctionInvocation;
import io.substrait.expression.Expression;
import io.substrait.type.Type;
import io.substrait.type.TypeCreator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.immutables.value.Value;

@Value.Immutable
public abstract class Aggregate extends SingleInputRel implements HasExtension {
  static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(Aggregate.class);

  public abstract List getGroupings();

  public abstract List getMeasures();

  @Override
  protected Type.Struct deriveRecordType() {
    return TypeCreator.REQUIRED.struct(
        Stream.concat(
            // unique grouping expressions
            getGroupings().stream()
                .flatMap(g -> g.getExpressions().stream())
                .collect(Collectors.toCollection(LinkedHashSet::new))
                .stream()
                .map(Expression::getType),

            // measures
            getMeasures().stream().map(t -> t.getFunction().getType())));
  }

  @Override
  public  O accept(RelVisitor visitor) throws E {
    return visitor.visit(this);
  }

  @Value.Immutable
  public abstract static class Grouping {
    public abstract List getExpressions();

    public static ImmutableGrouping.Builder builder() {
      return ImmutableGrouping.builder();
    }
  }

  @Value.Immutable
  public abstract static class Measure {
    public abstract AggregateFunctionInvocation getFunction();

    public abstract Optional getPreMeasureFilter();

    public static ImmutableMeasure.Builder builder() {
      return ImmutableMeasure.builder();
    }
  }

  public static ImmutableAggregate.Builder builder() {
    return ImmutableAggregate.builder();
  }
}