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.asterix.optimizer.rules;
  
  import org.apache.asterix.om.base.AInt32;
  import org.apache.asterix.om.base.IAObject;
  import org.apache.asterix.om.types.ATypeTag;
  import org.apache.asterix.om.utils.ConstantExpressionUtil;
  import org.apache.commons.lang3.mutable.Mutable;
  import org.apache.hyracks.algebricks.common.exceptions.AlgebricksException;
  import org.apache.hyracks.algebricks.core.algebra.base.ILogicalExpression;
  import org.apache.hyracks.algebricks.core.algebra.base.ILogicalOperator;
  import org.apache.hyracks.algebricks.core.algebra.base.IOptimizationContext;
  import org.apache.hyracks.algebricks.core.algebra.base.LogicalOperatorTag;
  import org.apache.hyracks.algebricks.core.algebra.base.PhysicalOperatorTag;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.AbstractLogicalOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.LimitOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.logical.OrderOperator;
  import org.apache.hyracks.algebricks.core.algebra.operators.physical.StableSortPOperator;
  import org.apache.hyracks.algebricks.core.algebra.util.OperatorPropertiesUtil;
  import org.apache.hyracks.algebricks.core.rewriter.base.IAlgebraicRewriteRule;
  
  /**
   * If an ORDER operator is followed by LIMIT, then we can push LIMIT into ORDER operator.
   * Finally, ORDER operator use TopKSorterOperatorDescriptor that can efficiently
   * sort tuples and fetch top K results.
   * =================
   * matching pattern:
   * limit <- order
   * =
   * producing pattern:
   * limit <- new order (topK applied)
   */
  public class PushLimitIntoOrderByRule implements IAlgebraicRewriteRule {
  
      @Override
      public boolean rewritePost(Mutable<ILogicalOperator> opRef, IOptimizationContext context) {
          return false;
      }
  
      @Override
      public boolean rewritePre(Mutable<ILogicalOperator> opRef, IOptimizationContext context)
              throws AlgebricksException {
          AbstractLogicalOperator op = (AbstractLogicalOperator) opRef.getValue();
          // The current operator should be LIMIT operator.
          if (op.getOperatorTag() != LogicalOperatorTag.LIMIT) {
              return false;
          }
  
          Mutable<ILogicalOperator> opRef2 = op.getInputs().get(0);
          AbstractLogicalOperator op2 = (AbstractLogicalOperator) opRef2.getValue();
  
          if (context.checkAndAddToAlreadyCompared(op, op2)) {
              return false;
          }
  
          // Should be ORDER operator
          if (op2.getOperatorTag() != LogicalOperatorTag.ORDER) {
              return false;
          } else {
              // ORDER operator is followed by LIMIT. Thus we can check whether we can apply this rule.
              boolean res = pushLimitIntoOrder(opRef, opRef2, context);
              if (res) {
                  OperatorPropertiesUtil.typeOpRec(opRef, context);
              }
              return res;
          }
      }
  
      /**
       * Generate new ORDER operator that uses TopKSort module and replaces the old ORDER operator.
       */
      private boolean pushLimitIntoOrder(Mutable<ILogicalOperator> opRef, Mutable<ILogicalOperator> opRef2,
              IOptimizationContext context) throws AlgebricksException {
          LimitOperator limitOp = (LimitOperator) opRef.getValue();
          OrderOperator orderOp = (OrderOperator) opRef2.getValue();
  
          // We don't push-down LIMIT into in-memory sort.
          if (orderOp.getPhysicalOperator().getOperatorTag() != PhysicalOperatorTag.STABLE_SORT) {
              return false;
          }
  
          Integer topK = getOutputLimit(limitOp);
          if (topK == null) {
              return false;
          }
  
          // Create the new ORDER operator, set the topK value, and replace the current one.
          OrderOperator newOrderOp = new OrderOperator(orderOp.getOrderExpressions(), topK);
          newOrderOp.setSourceLocation(orderOp.getSourceLocation());
          newOrderOp.setPhysicalOperator(new StableSortPOperator(newOrderOp.getTopK()));
          newOrderOp.getInputs().addAll(orderOp.getInputs());
          newOrderOp.setExecutionMode(orderOp.getExecutionMode());
          newOrderOp.recomputeSchema();
          newOrderOp.computeDeliveredPhysicalProperties(context);
          opRef2.setValue(newOrderOp);
          context.computeAndSetTypeEnvironmentForOperator(newOrderOp);
          context.addToDontApplySet(this, limitOp);
          return true;
      }
  
      static Integer getOutputLimit(LimitOperator limitOp) {
          if (!limitOp.hasMaxObjects()) {
              // No limit
              return null;
          }
          // Currently, we support LIMIT with a constant value.
          ILogicalExpression maxObjectsExpr = limitOp.getMaxObjects().getValue();
          IAObject maxObjectsValue = ConstantExpressionUtil.getConstantIaObject(maxObjectsExpr, ATypeTag.INTEGER);
          if (maxObjectsValue == null) {
              return null;
          }
          int topK = ((AInt32) maxObjectsValue).getIntegerValue();
          if (topK < 0) {
              topK = 0;
          }
  
          // Get the offset constant if there is one. If one presents, then topK = topK + offset.
          // This is because we can't apply offset to the external sort.
          // Final topK will be applied through LIMIT.
          if (limitOp.hasOffset()) {
              ILogicalExpression offsetExpr = limitOp.getOffset().getValue();
              IAObject offsetValue = ConstantExpressionUtil.getConstantIaObject(offsetExpr, ATypeTag.INTEGER);
              if (offsetValue == null) {
                  return null;
              }
              int offset = ((AInt32) offsetValue).getIntegerValue();
              if (offset < 0) {
                  offset = 0;
              }
              // Check the overflow case.
              if (offset >= Integer.MAX_VALUE - topK) {
                  return null;
              }
              topK += offset;
          }
  
          return topK;
      }
  }
  