com.hazelcast.shaded.org.apache.calcite.sql.SqlPivot 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.shaded.org.apache.calcite.sql;

import com.hazelcast.shaded.org.apache.calcite.sql.fun.SqlStdOperatorTable;
import com.hazelcast.shaded.org.apache.calcite.sql.parser.SqlParserPos;
import com.hazelcast.shaded.org.apache.calcite.sql.util.SqlBasicVisitor;
import com.hazelcast.shaded.org.apache.calcite.sql.util.SqlVisitor;
import com.hazelcast.shaded.org.apache.calcite.sql.validate.SqlValidatorUtil;
import com.hazelcast.shaded.org.apache.calcite.util.ImmutableNullableList;
import com.hazelcast.shaded.org.apache.calcite.util.Util;

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

import com.hazelcast.shaded.org.checkerframework.checker.nullness.qual.Nullable;

import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.stream.Collectors;

import static java.util.Objects.requireNonNull;

/**
 * Parse tree node that represents a PIVOT applied to a table reference
 * (or sub-query).
 *
 * Syntax:
 * 
{@code
 * SELECT *
 * FROM query PIVOT (agg, ... FOR axis, ... IN (in, ...)) AS alias}
 * 
 */
public class SqlPivot extends SqlCall {

  public SqlNode query;
  public final SqlNodeList aggList;
  public final SqlNodeList axisList;
  public final SqlNodeList inList;

  static final Operator OPERATOR = new Operator(SqlKind.PIVOT);

  //~ Constructors -----------------------------------------------------------

  public SqlPivot(SqlParserPos pos, SqlNode query, SqlNodeList aggList,
      SqlNodeList axisList, SqlNodeList inList) {
    super(pos);
    this.query = Objects.requireNonNull(query, "query");
    this.aggList = Objects.requireNonNull(aggList, "aggList");
    this.axisList = Objects.requireNonNull(axisList, "axisList");
    this.inList = Objects.requireNonNull(inList, "inList");
  }

  //~ Methods ----------------------------------------------------------------

  @Override public SqlOperator getOperator() {
    return OPERATOR;
  }

  @Override public List getOperandList() {
    return ImmutableNullableList.of(query, aggList, axisList, inList);
  }

  @SuppressWarnings("nullness")
  @Override public void setOperand(int i, @Nullable SqlNode operand) {
    // Only 'query' is mutable. (It is required for validation.)
    switch (i) {
    case 0:
      query = operand;
      break;
    default:
      super.setOperand(i, operand);
    }
  }

  @Override public void unparse(SqlWriter writer, int leftPrec, int rightPrec) {
    query.unparse(writer, leftPrec, 0);
    writer.keyword("PIVOT");
    final SqlWriter.Frame frame = writer.startList("(", ")");
    aggList.unparse(writer, 0, 0);
    writer.sep("FOR");
    // force parentheses if there is more than one axis
    final int leftPrec1 = axisList.size() > 1 ? 1 : 0;
    axisList.unparse(writer, leftPrec1, 0);
    writer.sep("IN");
    writer.list(SqlWriter.FrameTypeEnum.PARENTHESES, SqlWriter.COMMA,
        stripList(inList));
    writer.endList(frame);
  }

  static SqlNodeList stripList(SqlNodeList list) {
    return list.stream().map(SqlPivot::strip)
        .collect(SqlNode.toList(list.pos));
  }

  /** Converts a single-element SqlNodeList to its constituent node.
   * For example, "(1)" becomes "1";
   * "(2) as a" becomes "2 as a";
   * "(3, 4)" remains "(3, 4)";
   * "(5, 6) as b" remains "(5, 6) as b". */
  private static SqlNode strip(SqlNode e) {
    switch (e.getKind()) {
    case AS:
      final SqlCall call = (SqlCall) e;
      final List operands = call.getOperandList();
      return SqlStdOperatorTable.AS.createCall(e.pos,
          strip(operands.get(0)), operands.get(1));
    default:
      if (e instanceof SqlNodeList && ((SqlNodeList) e).size() == 1) {
        return ((SqlNodeList) e).get(0);
      }
      return e;
    }
  }

  /** Returns the aggregate list as (alias, call) pairs.
   * If there is no 'AS', alias is null. */
  public void forEachAgg(BiConsumer<@Nullable String, SqlNode> consumer) {
    for (SqlNode agg : aggList) {
      final SqlNode call = SqlUtil.stripAs(agg);
      final @Nullable String alias = SqlValidatorUtil.alias(agg);
      consumer.accept(alias, call);
    }
  }

  /** Returns the value list as (alias, node list) pairs. */
  public void forEachNameValues(BiConsumer consumer) {
    for (SqlNode node : inList) {
      String alias;
      if (node.getKind() == SqlKind.AS) {
        final List operands = ((SqlCall) node).getOperandList();
        alias = ((SqlIdentifier) operands.get(1)).getSimple();
        node = operands.get(0);
      } else {
        alias = pivotAlias(node);
      }
      consumer.accept(alias, toNodes(node));
    }
  }

  static String pivotAlias(SqlNode node) {
    if (node instanceof SqlNodeList) {
      return ((SqlNodeList) node).stream()
          .map(SqlPivot::pivotAlias).collect(Collectors.joining("_"));
    }
    return node.toString();
  }

  /** Converts a SqlNodeList to a list, and other nodes to a singleton list. */
  static SqlNodeList toNodes(SqlNode node) {
    if (node instanceof SqlNodeList) {
      return (SqlNodeList) node;
    } else {
      return new SqlNodeList(ImmutableList.of(node), node.getParserPosition());
    }
  }

  /** Returns the set of columns that are referenced as an argument to an
   * aggregate function or in a column in the {@code FOR} clause. All columns
   * that are not used will become "GROUP BY" columns. */
  public Set usedColumnNames() {
    final Set columnNames = new HashSet<>();
    final SqlVisitor nameCollector = new SqlBasicVisitor() {
      @Override public Void visit(SqlIdentifier id) {
        columnNames.add(Util.last(id.names));
        return super.visit(id);
      }
    };
    for (SqlNode agg : aggList) {
      final SqlCall call = (SqlCall) SqlUtil.stripAs(agg);
      call.accept(nameCollector);
    }
    for (SqlNode axis : axisList) {
      axis.accept(nameCollector);
    }
    return columnNames;
  }

  /** Pivot operator. */
  static class Operator extends SqlSpecialOperator {
    Operator(SqlKind kind) {
      super(kind.name(), kind);
    }

    @Override public SqlCall createCall(
        @Nullable SqlLiteral functionQualifier,
        SqlParserPos pos,
        @Nullable SqlNode... operands) {
      assert operands.length == 4;
      return new SqlPivot(pos, requireNonNull(operands[0], "query"),
          requireNonNull((SqlNodeList) operands[1], "aggList"),
          requireNonNull((SqlNodeList) operands[2], "axisList"),
          requireNonNull((SqlNodeList) operands[3], "inList"));
    }
  }
}