com.hazelcast.org.apache.calcite.adapter.enumerable.EnumerableTraitsUtils Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hazelcast-sql Show documentation
SQL Parser and Optimizer
There is a newer version: 5.5.0
/*
 * 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 com.hazelcast.org.apache.calcite.adapter.enumerable;

import com.hazelcast.org.apache.calcite.linq4j.Ord;
import com.hazelcast.org.apache.calcite.plan.RelOptUtil;
import com.hazelcast.org.apache.calcite.plan.RelTraitSet;
import com.hazelcast.org.apache.calcite.rel.RelCollation;
import com.hazelcast.org.apache.calcite.rel.RelCollations;
import com.hazelcast.org.apache.calcite.rel.RelFieldCollation;
import com.hazelcast.org.apache.calcite.rel.core.JoinRelType;
import com.hazelcast.org.apache.calcite.rel.type.RelDataType;
import com.hazelcast.org.apache.calcite.rel.type.RelDataTypeFactory;
import com.hazelcast.org.apache.calcite.rex.RexCall;
import com.hazelcast.org.apache.calcite.rex.RexCallBinding;
import com.hazelcast.org.apache.calcite.rex.RexInputRef;
import com.hazelcast.org.apache.calcite.rex.RexNode;
import com.hazelcast.org.apache.calcite.rex.RexUtil;
import com.hazelcast.org.apache.calcite.sql.SqlKind;
import com.hazelcast.org.apache.calcite.sql.validate.SqlMonotonicity;
import com.hazelcast.org.apache.calcite.util.Pair;
import com.hazelcast.org.apache.calcite.util.mapping.MappingType;
import com.hazelcast.org.apache.calcite.util.mapping.Mappings;

import com.hazelcast.com.google.common.collect.ImmutableList;

import org.apiguardian.api.API;
import com.hazelcast.org.checkerframework.checker.nullness.qual.Nullable;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * Utilities for traits propagation.
 */
@API(since = "1.24", status = API.Status.INTERNAL)
class EnumerableTraitsUtils {

  private EnumerableTraitsUtils() {}

  /**
   * Determine whether there is mapping between project input and output fields.
   * Bail out if sort relies on non-trivial expressions.
   */
  private static boolean isCollationOnTrivialExpr(
      List projects, RelDataTypeFactory typeFactory,
      Mappings.TargetMapping map, RelFieldCollation fc, boolean passDown) {
    final int index = fc.getFieldIndex();
    int target = map.getTargetOpt(index);
    if (target < 0) {
      return false;
    }

    final RexNode node = passDown ? projects.get(index) : projects.get(target);
    if (node.isA(SqlKind.CAST)) {
      // Check whether it is a monotonic preserving cast
      final RexCall cast = (RexCall) node;
      RelFieldCollation newFieldCollation = Objects.requireNonNull(RexUtil.apply(map, fc));
      final RexCallBinding binding =
          RexCallBinding.create(typeFactory, cast,
              ImmutableList.of(RelCollations.of(newFieldCollation)));
      if (cast.getOperator().getMonotonicity(binding)
          == SqlMonotonicity.NOT_MONOTONIC) {
        return false;
      }
    }

    return true;
  }

  static @Nullable Pair> passThroughTraitsForProject(
      RelTraitSet required,
      List exps,
      RelDataType inputRowType,
      RelDataTypeFactory typeFactory,
      RelTraitSet currentTraits) {
    final RelCollation collation = required.getCollation();
    if (collation == null || collation == RelCollations.EMPTY) {
      return null;
    }

    final Mappings.TargetMapping map =
        RelOptUtil.permutationIgnoreCast(
            exps, inputRowType);

    if (collation.getFieldCollations().stream().anyMatch(
        rc -> !isCollationOnTrivialExpr(exps, typeFactory,
            map, rc, true))) {
      return null;
    }

    final RelCollation newCollation = collation.apply(map);
    return Pair.of(currentTraits.replace(collation),
        ImmutableList.of(currentTraits.replace(newCollation)));
  }

  static @Nullable Pair> deriveTraitsForProject(
      RelTraitSet childTraits, int childId, List exps,
      RelDataType inputRowType, RelDataTypeFactory typeFactory, RelTraitSet currentTraits) {
    final RelCollation collation = childTraits.getCollation();
    if (collation == null || collation == RelCollations.EMPTY) {
      return null;
    }

    final int maxField = Math.max(exps.size(),
        inputRowType.getFieldCount());
    Mappings.TargetMapping mapping = Mappings
        .create(MappingType.FUNCTION, maxField, maxField);
    for (Ord node : Ord.zip(exps)) {
      if (node.e instanceof RexInputRef) {
        mapping.set(((RexInputRef) node.e).getIndex(), node.i);
      } else if (node.e.isA(SqlKind.CAST)) {
        final RexNode operand = ((RexCall) node.e).getOperands().get(0);
        if (operand instanceof RexInputRef) {
          mapping.set(((RexInputRef) operand).getIndex(), node.i);
        }
      }
    }

    List collationFieldsToDerive = new ArrayList<>();
    for (RelFieldCollation rc : collation.getFieldCollations()) {
      if (isCollationOnTrivialExpr(exps, typeFactory, mapping, rc, false)) {
        collationFieldsToDerive.add(rc);
      } else {
        break;
      }
    }

    if (collationFieldsToDerive.size() > 0) {
      final RelCollation newCollation = RelCollations
          .of(collationFieldsToDerive).apply(mapping);
      return Pair.of(currentTraits.replace(newCollation),
          ImmutableList.of(currentTraits.replace(collation)));
    } else {
      return null;
    }
  }

  /**
   * This function can be reused when a Join's traits pass-down shall only
   * pass through collation to left input.
   *
   * @param required required trait set for the join
   * @param joinType the join type
   * @param leftInputFieldCount number of field count of left join input
   * @param joinTraitSet trait set of the join
   */
  static @Nullable Pair> passThroughTraitsForJoin(
      RelTraitSet required, JoinRelType joinType,
      int leftInputFieldCount, RelTraitSet joinTraitSet) {
    RelCollation collation = required.getCollation();
    if (collation == null
        || collation == RelCollations.EMPTY
        || joinType == JoinRelType.FULL
        || joinType == JoinRelType.RIGHT) {
      return null;
    }

    for (RelFieldCollation fc : collation.getFieldCollations()) {
      // If field collation belongs to right input: cannot push down collation.
      if (fc.getFieldIndex() >= leftInputFieldCount) {
        return null;
      }
    }

    RelTraitSet passthroughTraitSet = joinTraitSet.replace(collation);
    return Pair.of(passthroughTraitSet,
        ImmutableList.of(
            passthroughTraitSet,
            passthroughTraitSet.replace(RelCollations.EMPTY)));
  }

  /**
   * This function can be reused when a Join's traits derivation shall only
   * derive collation from left input.
   *
   * @param childTraits trait set of the child
   * @param childId id of the child (0 is left join input)
   * @param joinType the join type
   * @param joinTraitSet trait set of the join
   * @param rightTraitSet trait set of the right join input
   */
  static @Nullable Pair> deriveTraitsForJoin(
      RelTraitSet childTraits, int childId, JoinRelType joinType,
      RelTraitSet joinTraitSet, RelTraitSet rightTraitSet) {
    // should only derive traits (limited to collation for now) from left join input.
    assert childId == 0;

    RelCollation collation = childTraits.getCollation();
    if (collation == null
        || collation == RelCollations.EMPTY
        || joinType == JoinRelType.FULL
        || joinType == JoinRelType.RIGHT) {
      return null;
    }

    RelTraitSet derivedTraits = joinTraitSet.replace(collation);
    return Pair.of(
        derivedTraits,
        ImmutableList.of(derivedTraits, rightTraitSet));
  }
}