exec/java-exec/src/main/java/org/apache/drill/exec/physical/impl/project/OutputWidthVisitor.java

/*
 * 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.drill.exec.physical.impl.project;

import com.google.common.base.Preconditions;
import org.apache.drill.common.expression.FunctionHolderExpression;
import org.apache.drill.common.expression.IfExpression;
import org.apache.drill.common.expression.LogicalExpression;
import org.apache.drill.common.expression.TypedNullConstant;
import org.apache.drill.common.expression.ValueExpressions;
import org.apache.drill.common.types.TypeProtos.MajorType;
import org.apache.drill.common.types.Types;
import org.apache.drill.exec.expr.AbstractExecExprVisitor;
import org.apache.drill.exec.expr.DrillFuncHolderExpr;
import org.apache.drill.exec.expr.ValueVectorReadExpression;
import org.apache.drill.exec.expr.ValueVectorWriteExpression;
import org.apache.drill.exec.expr.annotations.FunctionTemplate;
import org.apache.drill.exec.expr.fn.DrillFuncHolder;
import org.apache.drill.exec.expr.fn.output.OutputWidthCalculator;
import org.apache.drill.exec.physical.impl.project.OutputWidthExpression.FixedLenExpr;
import org.apache.drill.exec.physical.impl.project.OutputWidthExpression.FunctionCallExpr;
import org.apache.drill.exec.physical.impl.project.OutputWidthExpression.VarLenReadExpr;
import org.apache.drill.exec.physical.impl.project.OutputWidthExpression.IfElseWidthExpr;
import org.apache.drill.common.expression.ValueExpressions.VarDecimalExpression;
import org.apache.drill.exec.record.RecordBatchSizer;
import org.apache.drill.exec.record.TypedFieldId;

import java.util.ArrayList;

public class OutputWidthVisitor extends
      AbstractExecExprVisitor<OutputWidthExpression,
      OutputWidthVisitorState, RuntimeException> {

  @Override
  public OutputWidthExpression visitVarDecimalConstant(VarDecimalExpression varDecimalExpression,
                                                       OutputWidthVisitorState state) throws RuntimeException {
    Preconditions.checkArgument(varDecimalExpression.getMajorType().hasPrecision());
    return new FixedLenExpr(varDecimalExpression.getMajorType().getPrecision());
  }

  /**
   * Records the {@link IfExpression} as a {@link IfElseWidthExpr}.
   * IfElseWidthExpr will be reduced to a {@link FixedLenExpr} by taking the max
   * of the if-expr-width and the else-expr-width.
   */
  @Override
  public OutputWidthExpression visitIfExpression(IfExpression ifExpression, OutputWidthVisitorState state)
                                                                  throws RuntimeException {
    IfExpression.IfCondition condition = ifExpression.ifCondition;
    LogicalExpression ifExpr = condition.expression;
    LogicalExpression elseExpr = ifExpression.elseExpression;

    OutputWidthExpression ifWidthExpr = ifExpr.accept(this, state);
    OutputWidthExpression elseWidthExpr = null;
    if (elseExpr != null) {
      elseWidthExpr = elseExpr.accept(this, state);
    }
    return new IfElseWidthExpr(ifWidthExpr, elseWidthExpr);
  }

  /**
   * Handles a {@link FunctionHolderExpression}. Functions that produce
   * fixed-width output are trivially converted to a {@link FixedLenExpr}. For
   * functions that produce variable width output, the output width calculator
   * annotation is looked-up and recorded in a {@link FunctionCallExpr}. This
   * calculator will later be used to convert the FunctionCallExpr to a
   * {@link FixedLenExpr} expression
   */
  @Override
  public OutputWidthExpression visitFunctionHolderExpression(FunctionHolderExpression holderExpr,
                                                             OutputWidthVisitorState state) throws RuntimeException {
    OutputWidthExpression fixedWidth = getFixedLenExpr(holderExpr.getMajorType());
    if (fixedWidth != null) { return fixedWidth; }
    // Only Drill functions can be handled. Non-drill Functions, like HiveFunctions
    // will default to a fixed value
    if (!(holderExpr instanceof DrillFuncHolderExpr)) {
      // We currently only know how to handle DrillFuncs.
      // Use a default if this is not a DrillFunc
      return new FixedLenExpr(OutputSizeEstimateConstants.NON_DRILL_FUNCTION_OUTPUT_SIZE_ESTIMATE);
    }

    final DrillFuncHolder holder = ((DrillFuncHolderExpr) holderExpr).getHolder();

    // If the user has provided a size estimate, use it
    int estimate = holder.variableOutputSizeEstimate();
    if (estimate != FunctionTemplate.OUTPUT_SIZE_ESTIMATE_DEFAULT) {
      return new FixedLenExpr(estimate);
    }
    // Use the calculator provided by the user or use the default
    OutputWidthCalculator widthCalculator = holder.getOutputWidthCalculator();
    final int argSize = holderExpr.args.size();
    ArrayList<OutputWidthExpression> arguments = null;
    if (argSize != 0) {
      arguments = new ArrayList<>(argSize);
      for (LogicalExpression expr : holderExpr.args) {
        arguments.add(expr.accept(this, state));
      }
    }
    return new FunctionCallExpr(holderExpr, widthCalculator, arguments);
  }

  /**
   * Records a variable width write expression. This will be converted to a
   * {@link FixedLenExpr} expression by walking the tree of expression attached
   * to the write expression.
   */
  @Override
  public OutputWidthExpression visitValueVectorWriteExpression(ValueVectorWriteExpression writeExpr,
                                                               OutputWidthVisitorState state) throws RuntimeException {
    TypedFieldId fieldId = writeExpr.getFieldId();
    ProjectMemoryManager manager = state.getManager();
    OutputWidthExpression outputExpr;
    if (manager.isFixedWidth(fieldId)) {
      outputExpr = getFixedLenExpr(fieldId.getFinalType());
    } else {
      LogicalExpression writeArg = writeExpr.getChild();
      outputExpr = writeArg.accept(this, state);
    }
    return outputExpr;
  }

  /**
   * Records a variable width read expression as a {@link VarLenReadExpr}. This
   * will be converted to a {@link FixedLenExpr} expression by getting the size
   * for the corresponding column from the {@link RecordBatchSizer}.
   */
  @Override
  public OutputWidthExpression visitValueVectorReadExpression(ValueVectorReadExpression readExpr,
                                                              OutputWidthVisitorState state) throws RuntimeException {
    return new VarLenReadExpr(readExpr);
  }

  @Override
  public OutputWidthExpression visitQuotedStringConstant(ValueExpressions.QuotedString quotedString,
                                                         OutputWidthVisitorState state) throws RuntimeException {
    return new FixedLenExpr(quotedString.getString().length());
  }

  @Override
  public OutputWidthExpression visitUnknown(LogicalExpression logicalExpression, OutputWidthVisitorState state) {
    OutputWidthExpression fixedLenExpr = getFixedLenExpr(logicalExpression.getMajorType());
    if (fixedLenExpr != null) {
      return fixedLenExpr;
    }
    throw new IllegalStateException("Unknown variable width expression: " + logicalExpression);
  }

  @Override
  public OutputWidthExpression visitNullConstant(TypedNullConstant nullConstant, OutputWidthVisitorState state)
          throws RuntimeException {
    int width;
    if (nullConstant.getMajorType().hasPrecision()) {
      width = nullConstant.getMajorType().getPrecision();
    } else {
      width = 0;
    }
    return new FixedLenExpr(width);
  }


  @Override
  public OutputWidthExpression visitFixedLenExpr(FixedLenExpr fixedLenExpr, OutputWidthVisitorState state)
          throws RuntimeException {
    return fixedLenExpr;
  }

  /**
   * Converts the {@link VarLenReadExpr} to a {@link FixedLenExpr} by getting
   * the size for the corresponding column from the RecordBatchSizer.
   */
  @Override
  public OutputWidthExpression visitVarLenReadExpr(VarLenReadExpr varLenReadExpr, OutputWidthVisitorState state)
                                                      throws RuntimeException {
    String columnName = varLenReadExpr.getInputColumnName();
    if (columnName == null) {
      TypedFieldId fieldId = varLenReadExpr.getReadExpression().getTypedFieldId();
      columnName =  TypedFieldId.getPath(fieldId, state.manager.incomingBatch());
    }
    final RecordBatchSizer.ColumnSize columnSize = state.manager.getColumnSize(columnName);

    int columnWidth = columnSize.getDataSizePerEntry();
    return new FixedLenExpr(columnWidth);
  }

  /**
   * Converts a {@link FunctionCallExpr} to a {@link FixedLenExpr} by passing
   * the the args of the function to the width calculator for this function.
   */
  @Override
  public OutputWidthExpression visitFunctionCallExpr(FunctionCallExpr functionCallExpr, OutputWidthVisitorState state)
                                                      throws RuntimeException {
    ArrayList<OutputWidthExpression> args = functionCallExpr.getArgs();
    ArrayList<FixedLenExpr> estimatedArgs = null;

    if (args != null && args.size() != 0) {
      estimatedArgs = new ArrayList<>(args.size());
      for (OutputWidthExpression expr : args) {
        // Once the args are visited, they will all become FixedWidthExpr
        FixedLenExpr fixedLenExpr = (FixedLenExpr) expr.accept(this, state);
        estimatedArgs.add(fixedLenExpr);
      }
    }
    OutputWidthCalculator estimator = functionCallExpr.getCalculator();
    int estimatedSize = estimator.getOutputWidth(estimatedArgs);
    return new FixedLenExpr(estimatedSize);
  }

  /**
   * Converts the {@link IfElseWidthExpr} to a {@link FixedLenExpr} by taking
   * the max of the if-expr-width and the else-expr-width.
   */
  @Override
  public OutputWidthExpression visitIfElseWidthExpr(IfElseWidthExpr ifElseWidthExpr, OutputWidthVisitorState state)
                                                      throws RuntimeException {
    OutputWidthExpression ifReducedExpr = ifElseWidthExpr.expressions[0].accept(this, state);
    assert ifReducedExpr instanceof FixedLenExpr;
    int ifWidth = ((FixedLenExpr)ifReducedExpr).getDataWidth();
    int elseWidth = -1;
    if (ifElseWidthExpr.expressions[1] != null) {
      OutputWidthExpression elseReducedExpr = ifElseWidthExpr.expressions[1].accept(this, state);
      assert elseReducedExpr instanceof FixedLenExpr;
      elseWidth = ((FixedLenExpr)elseReducedExpr).getDataWidth();
    }
    int outputWidth = Math.max(ifWidth, elseWidth);
    return new FixedLenExpr(outputWidth);
  }

  private OutputWidthExpression getFixedLenExpr(MajorType majorType) {
    MajorType type = majorType;
    if (Types.isFixedWidthType(type)) {
      // Use only the width of the data. Metadata width will be accounted for at the end
      // This is to avoid using metadata size in intermediate calculations
      int fixedDataWidth = ProjectMemoryManager.getFixedWidth(type);
      return new OutputWidthExpression.FixedLenExpr(fixedDataWidth);
    }
    return null;
  }
}