content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to you under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.calcite.adapter.druid;
  
  import org.apache.calcite.config.CalciteConnectionConfig;
  import org.apache.calcite.plan.RelOptCluster;
  import org.apache.calcite.plan.RelOptPredicateList;
  import org.apache.calcite.plan.RelOptRule;
  import org.apache.calcite.plan.RelOptRuleCall;
  import org.apache.calcite.plan.RelOptUtil;
  import org.apache.calcite.plan.RelRule;
  import org.apache.calcite.rel.RelNode;
  import org.apache.calcite.rel.core.Aggregate;
  import org.apache.calcite.rel.core.AggregateCall;
  import org.apache.calcite.rel.core.Filter;
  import org.apache.calcite.rel.core.Project;
  import org.apache.calcite.rel.core.Sort;
  import org.apache.calcite.rel.logical.LogicalFilter;
  import org.apache.calcite.rel.rules.AggregateExtractProjectRule;
  import org.apache.calcite.rel.rules.AggregateFilterTransposeRule;
  import org.apache.calcite.rel.rules.FilterAggregateTransposeRule;
  import org.apache.calcite.rel.rules.FilterProjectTransposeRule;
  import org.apache.calcite.rel.rules.ProjectFilterTransposeRule;
  import org.apache.calcite.rel.rules.SortProjectTransposeRule;
  import org.apache.calcite.rel.type.RelDataType;
  import org.apache.calcite.rel.type.RelDataTypeFactory;
  import org.apache.calcite.rex.RexBuilder;
  import org.apache.calcite.rex.RexCall;
  import org.apache.calcite.rex.RexExecutor;
  import org.apache.calcite.rex.RexInputRef;
  import org.apache.calcite.rex.RexLiteral;
  import org.apache.calcite.rex.RexNode;
  import org.apache.calcite.rex.RexShuttle;
  import org.apache.calcite.rex.RexSimplify;
  import org.apache.calcite.rex.RexUtil;
  import org.apache.calcite.sql.SqlKind;
  import org.apache.calcite.sql.fun.SqlStdOperatorTable;
  import org.apache.calcite.tools.RelBuilder;
  import org.apache.calcite.util.Pair;
  import org.apache.calcite.util.Util;
  import org.apache.calcite.util.trace.CalciteTrace;
  
  import org.apache.commons.lang3.tuple.ImmutableTriple;
  import org.apache.commons.lang3.tuple.Triple;
  
  import com.google.common.collect.ImmutableList;
  import com.google.common.collect.Lists;
  
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.immutables.value.Value;
  import org.joda.time.Interval;
  import org.slf4j.Logger;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Set;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Rules and relational operators for {@link DruidQuery}.
   */
  public class DruidRules {
    private DruidRules() {}
  
    protected static final Logger LOGGER = CalciteTrace.getPlannerTracer();
  
    public static final DruidFilterRule FILTER =
        DruidFilterRule.DruidFilterRuleConfig.DEFAULT.toRule();
    public static final DruidProjectRule PROJECT =
        DruidProjectRule.DruidProjectRuleConfig.DEFAULT.toRule();
    public static final DruidAggregateRule AGGREGATE =
        DruidAggregateRule.DruidAggregateRuleConfig.DEFAULT.toRule();
    public static final DruidAggregateProjectRule AGGREGATE_PROJECT =
        DruidAggregateProjectRule.DruidAggregateProjectRuleConfig.DEFAULT.toRule();
    public static final DruidSortRule SORT =
        DruidSortRule.DruidSortRuleConfig.DEFAULT.toRule();
  
    /** Rule to push an {@link org.apache.calcite.rel.core.Sort} through a
     * {@link org.apache.calcite.rel.core.Project}. Useful to transform
     * to complex Druid queries. */
    public static final SortProjectTransposeRule SORT_PROJECT_TRANSPOSE =
        (SortProjectTransposeRule) SortProjectTransposeRule.Config.DEFAULT
            .withOperandFor(Sort.class, Project.class, DruidQuery.class)
            .withDescription("DruidSortProjectTransposeRule")
            .toRule();
  
    /** Rule to push a {@link org.apache.calcite.rel.core.Project}
     * past a {@link org.apache.calcite.rel.core.Filter}
     * when {@code Filter} is on top of a {@link DruidQuery}. */
    public static final ProjectFilterTransposeRule PROJECT_FILTER_TRANSPOSE =
        (ProjectFilterTransposeRule) ProjectFilterTransposeRule.Config.DEFAULT
            .withOperandFor(Project.class, Filter.class, DruidQuery.class)
            .withDescription("DruidProjectFilterTransposeRule")
            .toRule();
  
    /** Rule to push a {@link org.apache.calcite.rel.core.Filter}
     * past a {@link org.apache.calcite.rel.core.Project}
     * when {@code Project} is on top of a {@link DruidQuery}. */
    public static final FilterProjectTransposeRule FILTER_PROJECT_TRANSPOSE =
        (FilterProjectTransposeRule) FilterProjectTransposeRule.Config.DEFAULT
            .withOperandFor(Filter.class, Project.class, DruidQuery.class)
            .withCopyFilter(true)
            .withCopyProject(true)
            .withDescription("DruidFilterProjectTransposeRule")
            .toRule();
  
    /** Rule to push an {@link org.apache.calcite.rel.core.Aggregate}
     * past a {@link org.apache.calcite.rel.core.Filter}
     * when {@code Filter} is on top of a {@link DruidQuery}. */
    public static final AggregateFilterTransposeRule AGGREGATE_FILTER_TRANSPOSE =
        (AggregateFilterTransposeRule) AggregateFilterTransposeRule.Config.DEFAULT
            .withOperandFor(Aggregate.class, Filter.class, DruidQuery.class)
            .withDescription("DruidAggregateFilterTransposeRule")
            .toRule();
  
    /** Rule to push an {@link org.apache.calcite.rel.core.Filter}
     * past an {@link org.apache.calcite.rel.core.Aggregate}
     * when {@code Aggregate} is on top of a {@link DruidQuery}. */
    public static final FilterAggregateTransposeRule FILTER_AGGREGATE_TRANSPOSE =
        (FilterAggregateTransposeRule) FilterAggregateTransposeRule.Config.DEFAULT
            .withOperandFor(Filter.class, Aggregate.class, DruidQuery.class)
            .withDescription("DruidFilterAggregateTransposeRule")
            .toRule();
  
    public static final DruidPostAggregationProjectRule POST_AGGREGATION_PROJECT =
        DruidPostAggregationProjectRule.DruidPostAggregationProjectRuleConfig.DEFAULT.toRule();
  
    /** Rule to extract a {@link org.apache.calcite.rel.core.Project} from
     * {@link org.apache.calcite.rel.core.Aggregate} on top of
     * {@link org.apache.calcite.adapter.druid.DruidQuery} based on the fields
     * used in the aggregate. */
    public static final AggregateExtractProjectRule PROJECT_EXTRACT_RULE =
        (AggregateExtractProjectRule) AggregateExtractProjectRule.Config.DEFAULT
            .withOperandFor(Aggregate.class, DruidQuery.class)
            .withDescription("DruidAggregateExtractProjectRule")
            .toRule();
  
    public static final DruidHavingFilterRule DRUID_HAVING_FILTER_RULE =
        DruidHavingFilterRule.DruidHavingFilterRuleConfig.DEFAULT
        .toRule();
  
    public static final List<RelOptRule> RULES =
        ImmutableList.of(FILTER,
            PROJECT_FILTER_TRANSPOSE,
            AGGREGATE_FILTER_TRANSPOSE,
            AGGREGATE_PROJECT,
            PROJECT_EXTRACT_RULE,
            PROJECT,
            POST_AGGREGATION_PROJECT,
            AGGREGATE,
            FILTER_AGGREGATE_TRANSPOSE,
            FILTER_PROJECT_TRANSPOSE,
            SORT,
            SORT_PROJECT_TRANSPOSE,
            DRUID_HAVING_FILTER_RULE);
  
    /**
     * Rule to push a {@link org.apache.calcite.rel.core.Filter} into a
     * {@link DruidQuery}.
     */
    public static class DruidFilterRule
        extends RelRule<DruidFilterRule.DruidFilterRuleConfig> {
  
      /** Creates a DruidFilterRule. */
      protected DruidFilterRule(DruidFilterRuleConfig config) {
        super(config);
      }
  
      @Override public void onMatch(RelOptRuleCall call) {
        final Filter filter = call.rel(0);
        final DruidQuery query = call.rel(1);
        final RelOptCluster cluster = filter.getCluster();
        final RelBuilder relBuilder = call.builder();
        final RexBuilder rexBuilder = cluster.getRexBuilder();
  
        if (!DruidQuery.isValidSignature(query.signature() + 'f')) {
          return;
        }
  
        final List<RexNode> validPreds = new ArrayList<>();
        final List<RexNode> nonValidPreds = new ArrayList<>();
        final RexExecutor executor =
            Util.first(cluster.getPlanner().getExecutor(), RexUtil.EXECUTOR);
        final RelOptPredicateList predicates =
            call.getMetadataQuery().getPulledUpPredicates(filter.getInput());
        final RexSimplify simplify =
            new RexSimplify(rexBuilder, predicates, executor);
        final RexNode cond =
            simplify.simplifyUnknownAsFalse(filter.getCondition());
        for (RexNode e : RelOptUtil.conjunctions(cond)) {
          DruidJsonFilter druidJsonFilter =
              DruidJsonFilter.toDruidFilters(e, filter.getInput().getRowType(),
                  query, rexBuilder);
          if (druidJsonFilter != null) {
            validPreds.add(e);
          } else {
            nonValidPreds.add(e);
          }
        }
  
        // Timestamp
        int timestampFieldIdx =
            query.getRowType().getFieldNames()
                .indexOf(query.druidTable.timestampFieldName);
        RelNode newDruidQuery = query;
        final Triple<List<RexNode>, List<RexNode>, List<RexNode>> triple =
            splitFilters(validPreds, nonValidPreds, timestampFieldIdx);
        if (triple.getLeft().isEmpty() && triple.getMiddle().isEmpty()) {
          // it sucks, nothing to push
          return;
        }
        final List<RexNode> residualPreds = new ArrayList<>(triple.getRight());
        List<Interval> intervals = null;
        if (!triple.getLeft().isEmpty()) {
          final CalciteConnectionConfig connectionConfig =
              requireNonNull(
                  cluster.getPlanner().getContext()
                      .unwrap(CalciteConnectionConfig.class));
          requireNonNull(connectionConfig.timeZone());
          intervals =
              DruidDateTimeUtils.createInterval(
                  RexUtil.composeConjunction(rexBuilder, triple.getLeft()));
          if (intervals == null || intervals.isEmpty()) {
            // Case we have a filter with extract that can not be written as interval push down
            triple.getMiddle().addAll(triple.getLeft());
          }
        }
  
        if (!triple.getMiddle().isEmpty()) {
          final RelNode newFilter =
              filter.copy(filter.getTraitSet(), Util.last(query.rels),
                  RexUtil.composeConjunction(rexBuilder, triple.getMiddle()));
          newDruidQuery = DruidQuery.extendQuery(query, newFilter);
        }
        if (intervals != null && !intervals.isEmpty()) {
          newDruidQuery = DruidQuery.extendQuery((DruidQuery) newDruidQuery, intervals);
        }
        if (!residualPreds.isEmpty()) {
          newDruidQuery = relBuilder
              .push(newDruidQuery)
              .filter(residualPreds)
              .build();
        }
        call.transformTo(newDruidQuery);
      }
  
      /**
       * Given a list of conditions that contain Druid valid operations and
       * a list that contains those that contain any non-supported operation,
       * it outputs a triple with three different categories:
       * 1-l) condition filters on the timestamp column,
       * 2-m) condition filters that can be pushed to Druid,
       * 3-r) condition filters that cannot be pushed to Druid.
       */
      private static Triple<List<RexNode>, List<RexNode>, List<RexNode>> splitFilters(
          final List<RexNode> validPreds,
          final List<RexNode> nonValidPreds, final int timestampFieldIdx) {
        final List<RexNode> timeRangeNodes = new ArrayList<>();
        final List<RexNode> pushableNodes = new ArrayList<>();
        final List<RexNode> nonPushableNodes = new ArrayList<>(nonValidPreds);
        // Number of columns with the dimensions and timestamp
        for (RexNode conj : validPreds) {
          final RelOptUtil.InputReferencedVisitor visitor = new RelOptUtil.InputReferencedVisitor();
          conj.accept(visitor);
          if (visitor.inputPosReferenced.contains(timestampFieldIdx)
              && visitor.inputPosReferenced.size() == 1) {
            timeRangeNodes.add(conj);
          } else {
            pushableNodes.add(conj);
          }
        }
        return ImmutableTriple.of(timeRangeNodes, pushableNodes, nonPushableNodes);
      }
  
      /** Rule configuration. */
      @Value.Immutable(singleton = false)
      public interface DruidFilterRuleConfig extends RelRule.Config {
        DruidFilterRuleConfig DEFAULT = ImmutableDruidFilterRuleConfig.builder()
            .withOperandSupplier(b0 ->
                b0.operand(Filter.class).oneInput(b1 ->
                    b1.operand(DruidQuery.class).noInputs()))
            .build();
  
        @Override default DruidFilterRule toRule() {
          return new DruidFilterRule(this);
        }
      }
    }
  
    /** Rule to Push a Having {@link Filter} into a {@link DruidQuery}. */
    public static class DruidHavingFilterRule
        extends RelRule<DruidHavingFilterRule.DruidHavingFilterRuleConfig> {
  
      /** Creates a DruidHavingFilterRule. */
      protected DruidHavingFilterRule(DruidHavingFilterRuleConfig config) {
        super(config);
      }
  
      @Override public void onMatch(RelOptRuleCall call) {
        final Filter filter = call.rel(0);
        final DruidQuery query = call.rel(1);
        final RelOptCluster cluster = filter.getCluster();
        final RexBuilder rexBuilder = cluster.getRexBuilder();
  
        if (!DruidQuery.isValidSignature(query.signature() + 'h')) {
          return;
        }
  
        final RexNode cond = filter.getCondition();
        final DruidJsonFilter druidJsonFilter =
            DruidJsonFilter.toDruidFilters(cond, query.getTopNode().getRowType(),
                query, rexBuilder);
        if (druidJsonFilter != null) {
          final RelNode newFilter = filter
              .copy(filter.getTraitSet(), Util.last(query.rels), filter.getCondition());
          final DruidQuery newDruidQuery = DruidQuery.extendQuery(query, newFilter);
          call.transformTo(newDruidQuery);
        }
      }
  
      /** Rule configuration. */
      @Value.Immutable(singleton = false)
      public interface DruidHavingFilterRuleConfig extends RelRule.Config {
        DruidHavingFilterRuleConfig DEFAULT = ImmutableDruidHavingFilterRuleConfig.builder()
            .withOperandSupplier(b0 ->
                b0.operand(Filter.class).oneInput(b1 ->
                    b1.operand(DruidQuery.class).noInputs()))
            .build();
  
        @Override default DruidHavingFilterRule toRule() {
          return new DruidHavingFilterRule(this);
        }
      }
    }
  
    /**
     * Rule to push a {@link org.apache.calcite.rel.core.Project} into a
     * {@link DruidQuery}.
     */
    public static class DruidProjectRule
        extends RelRule<DruidProjectRule.DruidProjectRuleConfig> {
  
      /** Creates a DruidProjectRule. */
      protected DruidProjectRule(DruidProjectRuleConfig config) {
        super(config);
      }
  
      @Override public boolean matches(RelOptRuleCall call) {
        final Project project = call.rel(0);
        return project.getVariablesSet().isEmpty();
      }
  
      @Override public void onMatch(RelOptRuleCall call) {
        final Project project = call.rel(0);
        final DruidQuery query = call.rel(1);
        final RelOptCluster cluster = project.getCluster();
        final RexBuilder rexBuilder = cluster.getRexBuilder();
        if (!DruidQuery.isValidSignature(query.signature() + 'p')) {
          return;
        }
  
        if (DruidQuery.computeProjectAsScan(project, query.getTable().getRowType(), query)
            != null) {
          // All expressions can be pushed to Druid in their entirety.
          final RelNode newProject =
              project.copy(project.getTraitSet(),
                  ImmutableList.of(Util.last(query.rels)));
          RelNode newNode = DruidQuery.extendQuery(query, newProject);
          call.transformTo(newNode);
          return;
        }
  
        final Pair<List<RexNode>, List<RexNode>> pair =
            splitProjects(rexBuilder, query, project.getProjects());
        if (pair == null) {
          // We can't push anything useful to Druid.
          return;
        }
        final List<RexNode> above = pair.left;
        final List<RexNode> below = pair.right;
        final RelDataTypeFactory.Builder builder =
            cluster.getTypeFactory().builder();
        final RelNode input = Util.last(query.rels);
        for (RexNode e : below) {
          final String name;
          if (e instanceof RexInputRef) {
            name = input.getRowType().getFieldNames().get(((RexInputRef) e).getIndex());
          } else {
            name = null;
          }
          builder.add(name, e.getType());
        }
        final RelNode newProject =
            project.copy(project.getTraitSet(), input, below, builder.build());
        final DruidQuery newQuery = DruidQuery.extendQuery(query, newProject);
        final RelNode newProject2 =
                project.copy(project.getTraitSet(), newQuery, above,
                    project.getRowType());
        call.transformTo(newProject2);
      }
  
      private static @Nullable Pair<List<RexNode>, List<RexNode>> splitProjects(
          final RexBuilder rexBuilder, final RelNode input, List<RexNode> nodes) {
        final RelOptUtil.InputReferencedVisitor visitor =
            new RelOptUtil.InputReferencedVisitor();
        visitor.visitEach(nodes);
        if (visitor.inputPosReferenced.size() == input.getRowType().getFieldCount()) {
          // All inputs are referenced
          return null;
        }
        final List<RexNode> belowNodes = new ArrayList<>();
        final List<RelDataType> belowTypes = new ArrayList<>();
        final List<Integer> positions = Lists.newArrayList(visitor.inputPosReferenced);
        for (int i : positions) {
          final RexNode node = rexBuilder.makeInputRef(input, i);
          belowNodes.add(node);
          belowTypes.add(node.getType());
        }
        final List<RexNode> aboveNodes = new RexShuttle() {
          @Override public RexNode visitInputRef(RexInputRef ref) {
            final int index = positions.indexOf(ref.getIndex());
            return rexBuilder.makeInputRef(belowTypes.get(index), index);
          }
        }.visitList(nodes);
        return Pair.of(aboveNodes, belowNodes);
      }
  
      /** Rule configuration. */
      @Value.Immutable(singleton = false)
      public interface DruidProjectRuleConfig extends RelRule.Config {
        DruidProjectRuleConfig DEFAULT = ImmutableDruidProjectRuleConfig.builder()
            .withOperandSupplier(b0 ->
                b0.operand(Project.class).oneInput(b1 ->
                    b1.operand(DruidQuery.class).noInputs()))
            .build();
  
        @Override default DruidProjectRule toRule() {
          return new DruidProjectRule(this);
        }
      }
    }
  
    /**
     * Rule to push a {@link org.apache.calcite.rel.core.Project} into a
     * {@link DruidQuery} as a Post aggregator.
     */
    public static class DruidPostAggregationProjectRule
        extends RelRule<DruidPostAggregationProjectRule.DruidPostAggregationProjectRuleConfig> {
  
      /** Creates a DruidPostAggregationProjectRule. */
      protected DruidPostAggregationProjectRule(DruidPostAggregationProjectRuleConfig config) {
        super(config);
      }
  
      @Override public void onMatch(RelOptRuleCall call) {
        Project project = call.rel(0);
        DruidQuery query = call.rel(1);
        if (!DruidQuery.isValidSignature(query.signature() + 'o')) {
          return;
        }
        boolean hasRexCalls = false;
        for (RexNode rexNode : project.getProjects()) {
          if (rexNode instanceof RexCall) {
            hasRexCalls = true;
            break;
          }
        }
        // Only try to push down Project when there will be Post aggregators in result DruidQuery
        if (hasRexCalls) {
  
          final RelNode topNode = query.getTopNode();
          final Aggregate topAgg;
          if (topNode instanceof Aggregate) {
            topAgg = (Aggregate) topNode;
          } else {
            topAgg = (Aggregate) ((Filter) topNode).getInput();
          }
  
          for (RexNode rexNode : project.getProjects()) {
            if (DruidExpressions.toDruidExpression(rexNode, topAgg.getRowType(), query) == null) {
              return;
            }
          }
          final RelNode newProject = project
              .copy(project.getTraitSet(), ImmutableList.of(Util.last(query.rels)));
          final DruidQuery newQuery = DruidQuery.extendQuery(query, newProject);
          call.transformTo(newQuery);
        }
      }
  
      /** Rule configuration. */
      @Value.Immutable(singleton = false)
      public interface DruidPostAggregationProjectRuleConfig extends RelRule.Config {
        DruidPostAggregationProjectRuleConfig DEFAULT =
            ImmutableDruidPostAggregationProjectRuleConfig.builder()
                .withOperandSupplier(b0 ->
                    b0.operand(Project.class).oneInput(b1 ->
                        b1.operand(DruidQuery.class).noInputs()))
                .build();
  
        @Override default DruidPostAggregationProjectRule toRule() {
          return new DruidPostAggregationProjectRule(this);
        }
      }
    }
  
    /**
     * Rule to push an {@link org.apache.calcite.rel.core.Aggregate}
     * into a {@link DruidQuery}.
     */
    public static class DruidAggregateRule
        extends RelRule<DruidAggregateRule.DruidAggregateRuleConfig> {
  
      /** Creates a DruidAggregateRule. */
      protected DruidAggregateRule(DruidAggregateRuleConfig config) {
        super(config);
      }
  
      @Override public void onMatch(RelOptRuleCall call) {
        final Aggregate aggregate = call.rel(0);
        final DruidQuery query = call.rel(1);
        final RelNode topDruidNode = query.getTopNode();
        final Project project = topDruidNode instanceof Project ? (Project) topDruidNode : null;
        if (!DruidQuery.isValidSignature(query.signature() + 'a')) {
          return;
        }
  
        if (aggregate.getGroupSets().size() != 1) {
          return;
        }
        if (DruidQuery
            .computeProjectGroupSet(project, aggregate.getGroupSet(), query.table.getRowType(), query)
            == null) {
          return;
        }
        final List<String> aggNames = Util
            .skip(aggregate.getRowType().getFieldNames(), aggregate.getGroupSet().cardinality());
        if (DruidQuery.computeDruidJsonAgg(aggregate.getAggCallList(), aggNames, project, query)
            == null) {
          return;
        }
        final RelNode newAggregate = aggregate
            .copy(aggregate.getTraitSet(), ImmutableList.of(query.getTopNode()));
        call.transformTo(DruidQuery.extendQuery(query, newAggregate));
      }
  
      /** Rule configuration. */
      @Value.Immutable(singleton = false)
      public interface DruidAggregateRuleConfig extends RelRule.Config {
        DruidAggregateRuleConfig DEFAULT = ImmutableDruidAggregateRuleConfig.builder()
            .withOperandSupplier(b0 ->
                b0.operand(Aggregate.class).oneInput(b1 ->
                    b1.operand(DruidQuery.class).noInputs()))
            .build();
  
        @Override default DruidAggregateRule toRule() {
          return new DruidAggregateRule(this);
        }
      }
    }
  
    /**
     * Rule to push an {@link org.apache.calcite.rel.core.Aggregate} and
     * {@link org.apache.calcite.rel.core.Project} into a {@link DruidQuery}.
     */
    public static class DruidAggregateProjectRule
        extends RelRule<DruidAggregateProjectRule.DruidAggregateProjectRuleConfig> {
  
      /** Creates a DruidAggregateProjectRule. */
      protected DruidAggregateProjectRule(DruidAggregateProjectRuleConfig config) {
        super(config);
      }
  
      @Override public void onMatch(RelOptRuleCall call) {
        final Aggregate aggregate = call.rel(0);
        final Project project = call.rel(1);
        final DruidQuery query = call.rel(2);
        if (!DruidQuery.isValidSignature(query.signature() + 'p' + 'a')) {
          return;
        }
        if (aggregate.getGroupSets().size() != 1) {
          return;
        }
        if (DruidQuery
            .computeProjectGroupSet(project, aggregate.getGroupSet(), query.table.getRowType(), query)
            == null) {
          return;
        }
        final List<String> aggNames = Util
            .skip(aggregate.getRowType().getFieldNames(), aggregate.getGroupSet().cardinality());
        if (DruidQuery.computeDruidJsonAgg(aggregate.getAggCallList(), aggNames, project, query)
            == null) {
          return;
        }
        final RelNode newProject =
            project.copy(project.getTraitSet(),
                ImmutableList.of(Util.last(query.rels)));
        final RelNode newAggregate =
            aggregate.copy(aggregate.getTraitSet(),
                ImmutableList.of(newProject));
        List<Integer> filterRefs = getFilterRefs(aggregate.getAggCallList());
        final DruidQuery query2;
        if (!filterRefs.isEmpty()) {
          query2 =
              optimizeFilteredAggregations(call, query, (Project) newProject,
                  (Aggregate) newAggregate);
        } else {
          final DruidQuery query1 = DruidQuery.extendQuery(query, newProject);
          query2 = DruidQuery.extendQuery(query1, newAggregate);
        }
        call.transformTo(query2);
      }
  
      /** Returns an array of unique filter references from the given list of
       * {@link org.apache.calcite.rel.core.AggregateCall}s. */
      private static Set<Integer> getUniqueFilterRefs(List<AggregateCall> calls) {
        Set<Integer> refs = new HashSet<>();
        for (AggregateCall call : calls) {
          if (call.hasFilter()) {
            refs.add(call.filterArg);
          }
        }
        return refs;
      }
  
      /**
       * Attempts to optimize any aggregations with filters in the DruidQuery.
       * Uses the following steps:
       *
       * <ol>
       * <li>Tries to abstract common filters out into the "filter" field;
       * <li>Eliminates expressions that are always true or always false when
       *     possible;
       * <li>ANDs aggregate filters together with the outer filter to allow for
       *     pruning of data.
       * </ol>
       *
       * <p>Should be called before pushing both the aggregate and project into
       * Druid. Assumes that at least one aggregate call has a filter attached to
       * it. */
      private static DruidQuery optimizeFilteredAggregations(RelOptRuleCall call,
          DruidQuery query,
          Project project, Aggregate aggregate) {
        Filter filter = null;
        final RexBuilder builder = query.getCluster().getRexBuilder();
        final RexExecutor executor =
            Util.first(query.getCluster().getPlanner().getExecutor(),
                RexUtil.EXECUTOR);
        final RelNode scan = query.rels.get(0); // first rel is the table scan
        final RelOptPredicateList predicates =
            call.getMetadataQuery().getPulledUpPredicates(scan);
        final RexSimplify simplify =
            new RexSimplify(builder, predicates, executor);
  
        // if the druid query originally contained a filter
        boolean containsFilter = false;
        for (RelNode node : query.rels) {
          if (node instanceof Filter) {
            filter = (Filter) node;
            containsFilter = true;
            break;
          }
        }
  
        // if every aggregate call has a filter arg reference
        boolean allHaveFilters = allAggregatesHaveFilters(aggregate.getAggCallList());
  
        Set<Integer> uniqueFilterRefs = getUniqueFilterRefs(aggregate.getAggCallList());
  
        // One of the pre-conditions for this method
        assert !uniqueFilterRefs.isEmpty();
  
        List<AggregateCall> newCalls = new ArrayList<>();
  
        // OR all the filters so that they can ANDed to the outer filter
        List<RexNode> disjunctions = new ArrayList<>();
        for (Integer i : uniqueFilterRefs) {
          disjunctions.add(stripFilter(project.getProjects().get(i)));
        }
        RexNode filterNode = RexUtil.composeDisjunction(builder, disjunctions);
  
        // Erase references to filters
        for (AggregateCall aggCall : aggregate.getAggCallList()) {
          if ((uniqueFilterRefs.size() == 1
                  && allHaveFilters) // filters get extracted
              || aggCall.hasFilter()
              && project.getProjects().get(aggCall.filterArg).isAlwaysTrue()) {
            aggCall = aggCall.withFilter(-1);
          }
          newCalls.add(aggCall);
        }
        aggregate =
            aggregate.copy(aggregate.getTraitSet(), aggregate.getInput(),
                aggregate.getGroupSet(), aggregate.getGroupSets(), newCalls);
  
        if (containsFilter) {
          // AND the current filterNode with the filter node inside filter
          filterNode = builder.makeCall(SqlStdOperatorTable.AND, filterNode, filter.getCondition());
        }
  
        // Simplify the filter as much as possible
        RexNode tempFilterNode = filterNode;
        filterNode = simplify.simplifyUnknownAsFalse(filterNode);
  
        // It's possible that after simplification that the expression is now always false.
        // Druid cannot handle such a filter.
        // This will happen when the below expression (f_n+1 may not exist):
        // f_n+1 AND (f_1 OR f_2 OR ... OR f_n) simplifies to be something always false.
        // f_n+1 cannot be false, since it came from a pushed filter rel node
        // and each f_i cannot be false, since DruidAggregateProjectRule would have caught that.
        // So, the only solution is to revert back to the un simplified version and let Druid
        // handle a filter that is ultimately unsatisfiable.
        if (filterNode.isAlwaysFalse()) {
          filterNode = tempFilterNode;
        }
  
        filter = LogicalFilter.create(scan, filterNode);
  
        boolean addNewFilter = !filter.getCondition().isAlwaysTrue() && allHaveFilters;
        // Assumes that Filter nodes are always right after
        // TableScan nodes (which are always present)
        int startIndex = containsFilter && addNewFilter ? 2 : 1;
  
        List<RelNode> newNodes =
            constructNewNodes(query.rels, addNewFilter, startIndex,
                filter, project, aggregate);
  
        return DruidQuery.create(query.getCluster(),
               aggregate.getTraitSet().replace(query.getConvention()),
               query.getTable(), query.druidTable, newNodes);
      }
  
      // Returns true if and only if every AggregateCall in calls has a filter argument.
      private static boolean allAggregatesHaveFilters(List<AggregateCall> calls) {
        for (AggregateCall call : calls) {
          if (!call.hasFilter()) {
            return false;
          }
        }
        return true;
      }
  
      /**
       * Returns a new List of RelNodes in the order of the given order of the oldNodes,
       * the given {@link Filter}, and any extra nodes.
       */
      private static List<RelNode> constructNewNodes(List<RelNode> oldNodes,
          boolean addFilter, int startIndex, RelNode filter, RelNode... trailingNodes) {
        List<RelNode> newNodes = new ArrayList<>();
  
        // The first item should always be the Table scan, so any filter would go after that
        newNodes.add(oldNodes.get(0));
  
        if (addFilter) {
          newNodes.add(filter);
          // This is required so that each RelNode is linked to the one before it
          if (startIndex < oldNodes.size()) {
            RelNode next = oldNodes.get(startIndex);
            newNodes.add(next.copy(next.getTraitSet(), Collections.singletonList(filter)));
            startIndex++;
          }
        }
  
        // Add the rest of the nodes from oldNodes
        for (int i = startIndex; i < oldNodes.size(); i++) {
          newNodes.add(oldNodes.get(i));
        }
  
        // Add the trailing nodes (need to link them)
        for (RelNode node : trailingNodes) {
          newNodes.add(node.copy(node.getTraitSet(), Collections.singletonList(Util.last(newNodes))));
        }
  
        return newNodes;
      }
  
      // Removes the IS_TRUE in front of RexCalls, if they exist
      private static RexNode stripFilter(RexNode node) {
        if (node.getKind() == SqlKind.IS_TRUE) {
          return ((RexCall) node).getOperands().get(0);
        }
        return node;
      }
  
      private static List<Integer> getFilterRefs(List<AggregateCall> calls) {
        List<Integer> refs = new ArrayList<>();
        for (AggregateCall call : calls) {
          if (call.hasFilter()) {
            refs.add(call.filterArg);
          }
        }
        return refs;
      }
  
      /** Rule configuration. */
      @Value.Immutable(singleton = false)
      public interface DruidAggregateProjectRuleConfig extends RelRule.Config {
        DruidAggregateProjectRuleConfig DEFAULT = ImmutableDruidAggregateProjectRuleConfig.builder()
            .withOperandSupplier(b0 ->
                b0.operand(Aggregate.class).oneInput(b1 ->
                    b1.operand(Project.class).oneInput(b2 ->
                        b2.operand(DruidQuery.class).noInputs())))
            .build();
  
        @Override default DruidAggregateProjectRule toRule() {
          return new DruidAggregateProjectRule(this);
        }
      }
    }
  
    /**
     * Rule to push a {@link org.apache.calcite.rel.core.Sort}
     * into a {@link DruidQuery}.
     */
    public static class DruidSortRule
        extends RelRule<DruidSortRule.DruidSortRuleConfig> {
  
      /** Creates a DruidSortRule. */
      protected DruidSortRule(DruidSortRuleConfig config) {
        super(config);
      }
  
      @Override public void onMatch(RelOptRuleCall call) {
        final Sort sort = call.rel(0);
        final DruidQuery query = call.rel(1);
        if (!DruidQuery.isValidSignature(query.signature() + 'l')) {
          return;
        }
        // Either it is:
        // - a pure limit above a query of type scan
        // - a sort and limit on a dimension/metric part of the druid group by query
        if (sort.offset != null && RexLiteral.intValue(sort.offset) != 0) {
          // offset not supported by Druid
          return;
        }
        if (query.getQueryType() == QueryType.SCAN && !RelOptUtil.isPureLimit(sort)) {
          return;
        }
  
        final RelNode newSort = sort
            .copy(sort.getTraitSet(), ImmutableList.of(Util.last(query.rels)));
        call.transformTo(DruidQuery.extendQuery(query, newSort));
      }
  
      /** Rule configuration. */
      @Value.Immutable(singleton = false)
      public interface DruidSortRuleConfig extends RelRule.Config {
        DruidSortRuleConfig DEFAULT = ImmutableDruidSortRuleConfig.builder()
            .withOperandSupplier(b0 ->
                b0.operand(Sort.class).oneInput(b1 ->
                    b1.operand(DruidQuery.class).noInputs()))
            .build();
  
        @Override default DruidSortRule toRule() {
          return new DruidSortRule(this);
        }
      }
    }
  }
  