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.lang.sqlpp.parser;
  
  import static org.apache.asterix.lang.sqlpp.parser.SqlppGroupingSetsParser.GROUPING_SETS_LIMIT;
  
  import java.io.PrintWriter;
  import java.io.StringWriter;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.List;
  
  import org.apache.asterix.common.exceptions.CompilationException;
  import org.apache.asterix.common.exceptions.ErrorCode;
  import org.apache.asterix.common.metadata.NamespaceResolver;
  import org.apache.asterix.lang.common.base.Expression;
  import org.apache.asterix.lang.common.base.ILangExpression;
  import org.apache.asterix.lang.common.base.IParser;
  import org.apache.asterix.lang.common.base.Statement;
  import org.apache.asterix.lang.common.clause.GroupbyClause;
  import org.apache.asterix.lang.common.expression.GbyVariableExpressionPair;
  import org.apache.asterix.lang.sqlpp.util.SqlppVariableUtil;
  import org.apache.asterix.lang.sqlpp.visitor.SqlppFormatPrintVisitor;
  import org.apache.asterix.lang.sqlpp.visitor.base.AbstractSqlppSimpleExpressionVisitor;
  import org.apache.hyracks.api.exceptions.IError;
  import org.apache.hyracks.util.MathUtil;
  import org.junit.Assert;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.junit.runners.Parameterized;
  
  /**
   * Test rewriting from GROUP BY's grouping elements into grouping sets:
   * <p>
   * GROUP BY {0} => GROUP BY GROUPING SETS ({1})
   */
  @RunWith(Parameterized.class)
  public class SqlppGroupingSetsParserTest {
  
      private static final int GROUPING_SETS_LIMIT_LOG2 = MathUtil.log2Floor(GROUPING_SETS_LIMIT);
  
      private static final int GROUPING_SETS_LIMIT_SQRT = (int) Math.ceil(Math.sqrt(GROUPING_SETS_LIMIT));
  
      private static final ErrorCode ERR_SYNTAX = ErrorCode.PARSE_ERROR;
  
      private static final ErrorCode ERR_OVERFLOW = ErrorCode.COMPILATION_GROUPING_SETS_OVERFLOW;
  
      private static final ErrorCode ERR_ALIAS = ErrorCode.COMPILATION_UNEXPECTED_ALIAS;
  
      @Parameterized.Parameters(name = "{index}: GROUP BY {0}")
      public static Collection<Object[]> data() {
          return Arrays
                  .asList(new Object[][] {
                          // GROUP BY {0} => GROUP BY GROUPING SETS ({1})
                          //
                          // 1. Basic
                          //
                          { "()", "()" },
                          //
                          { "(), ()", "()" },
                          //
                          { "a", "(a)" },
                          //
                          { "a,b", "(a,b)" },
                          //
                          { "a,(b,c)", "(a,b,c)" },
                          //
                          { "(a,b),(c,d)", "(a,b,c,d)" },
                          //
                          // 2. Rollup
                          //
                          { "rollup(a,b,c)", "(a,b,c)(a,b)(a)()" },
                          //
                          { "Rollup(a,(b,c),d)", "(a,b,c,d)(a,b,c)(a)()" },
                          //
                          { "RollUP((a,b),(c,d))", "(a,b,c,d)(a,b)()" },
                          //
                          { "a,ROLLUP(a,b)", "(a,b)(a)(a)" },
                          //
                          { "a,ROLLUP(a,b),c", "(a,b,c)(a,c)(a,c)" },
                          //
                          { "a,b,ROLLUP(c,d)", "(a,b,c,d)(a,b,c)(a,b)" },
                          //
                          { "ROLLUP(a,b),ROLLUP(c,d)", "(a,b,c,d)(a,b,c)(a,b)(a,c,d)(a,c)(a)(c,d)(c)()" },
                          //
                          // 3. Cube
                          //
                          { "cube(a,b,c)", "(a,b,c)(a,b)(a,c)(a)(b,c)(b)(c)()" },
                          //
                          { "Cube(a,(b,c),d)", "(a,b,c,d)(a,b,c)(a,d)(a)(b,c,d)(b,c)(d)()" },
                          //
                          { "CubE((a,b),(c,d))", "(a,b,c,d)(a,b)(c,d)()" },
                          //
                          { "a,CUBE(a,b)", "(a,b)(a)(a,b)(a)" },
                          //
                          { "a,CUBE(a,b),c", "(a,b,c)(a,c)(a,b,c)(a,c)" },
                          //
                          { "a,b,CUBE(c,d)", "(a,b,c,d)(a,b,c)(a,b,d)(a,b)" },
                          //
                          { "CUBE(a,b),CUBE(c,d)",
                                  "(a,b,c,d)(a,b,c)(a,b,d)(a,b)(a,c,d)(a,c)(a,d)(a)(b,c,d)(b,c)(b,d)(b)(c,d)(c)(d)()" },
                          //
                          // 4. Rollup + Cube
                          //
                          { "ROLLUP(a,b),CUBE(c,d)", "(a,b,c,d)(a,b,c)(a,b,d)(a,b)(a,c,d)(a,c)(a,d)(a)(c,d)(c)(d)()" },
                          //
                          { "CUBE(a,b),ROLLUP(c,d)", "(a,b,c,d)(a,b,c)(a,b)(a,c,d)(a,c)(a)(b,c,d)(b,c)(b)(c,d)(c)()" },
                          //
                          // 5. Grouping Sets
                          //
                          { "grouping sets(())", "()" },
                          //
                          { "Grouping Sets((), ())", "()()" },
                          //
                          { "Grouping setS(()), GROUPING SETS(())", "()" },
                          //
                          { "GROUPING SETS((a),(a,b))", "(a)(a,b)" },
                          //
                          { "GROUPING SETS((a,b),(a,b))", "(a,b)(a,b)" },
                          //
                          { "GROUPING SETS((a,b)),GROUPING SETS((a,b))", "(a,b)" },
                          //
                          { "GROUPING SETS((a,b),(c)),GROUPING SETS((d,e),())", "(a,b,d,e)(a,b)(c,d,e)(c)" },
                          //
                          { "GROUPING SETS(ROLLUP(a,b),ROLLUP(c,d))", "(a,b)(a)()(c,d)(c)()" },
                          //
                          { "GROUPING SETS(ROLLUP(a,b)), GROUPING SETS(ROLLUP(c,d))",
                                  "(a,b,c,d)(a,b,c)(a,b)(a,c,d)(a,c)(a)(c,d)(c)()" },
                          //
                          { "GROUPING SETS((a, b), GROUPING SETS((c,d), (e,f)))", "(a,b)(c,d)(e,f)" },
                          //
                          { "GROUPING SETS(ROLLUP(a,b)),GROUPING SETS(ROLLUP(c,d))",
                                  "(a,b,c,d)(a,b,c)(a,b)(a,c,d)(a,c)(a)(c,d)(c)()" },
                          //
                          // 6. Variable names (aliases)
                          //
                          { "a as x, b as y", "(a as x,b as y)" },
                          //
                          { "ROLLUP(a as x, b as y),ROLLUP(a, b)",
                                  "(a as x,b as y)(a as x,b as y)(a as x,b as y)(a as x,b as y)"
                                          + "(a as x)(a as x)(a as x,b as y)(a as x)()" },
                          //
                          { "CUBE(a as x, b as y),CUBE(a, b)",
                                  "(a as x,b as y)(a as x,b as y)(a as x,b as y)(a as x,b as y)"
                                          + "(a as x,b as y)(a as x)(a as x,b as y)(a as x)(b as y,a as x)(b as y,a as x)(b as y)(b as y)"
                                          + "(a as x,b as y)(a as x)(b as y)()" },
                          //
                          { "GROUPING SETS((e1 as x, e2 as y)), GROUPING SETS((e1, e2))", "(e1 as x,e2 as y)" },
                          //
                          { "GROUPING SETS((e1 as x, e2 as y), (e1, e2))", "(e1 as x,e2 as y)(e1 as x,e2 as y)" },
                          //
                          // 7. Errors
                          //
                          // Syntax error
                          { "ROLLUP()", ERR_SYNTAX },
                          //
                          // Syntax error
                          { "CUBE()", ERR_SYNTAX },
                          //
                          // Too many grouping sets when expanding a rollup
                          { String.format("ROLLUP(%s)", generateSimpleGroupingSet("a", GROUPING_SETS_LIMIT)),
                                  ERR_OVERFLOW },
                          //
                          // Too many grouping sets when expanding a cube
                          { String.format("CUBE(%s)", generateSimpleGroupingSet("a", GROUPING_SETS_LIMIT_LOG2 + 1)),
                                  ERR_OVERFLOW },
                          //
                          // Too many grouping sets when doing a cross product of grouping sets
                          { String.format("GROUPING SETS(%s), GROUPING SETS(%s)",
                                  generateSimpleGroupingSet("a", GROUPING_SETS_LIMIT_SQRT),
                                  generateSimpleGroupingSet("b", GROUPING_SETS_LIMIT_SQRT)), ERR_OVERFLOW },
                          //
                          // Unexpected aliases
                          //
                          { "ROLLUP(a as x, b),ROLLUP(a, b as y)", ERR_ALIAS },
                          //
                          { "CUBE(a as x, b),CUBE(a, b as y)", ERR_ALIAS },
                          //
                          { "GROUPING SETS((e1 as x, e2), (e1, e2 as y))", ERR_ALIAS },
                          //
                          { "GROUPING SETS((e1 as x, e2)), GROUPING SETS((e1, e2 as y))", ERR_ALIAS },
                          //
                          { "GROUPING SETS((e1 as a, e2 as b)), GROUPING SETS((e1 as c, e2 as d))", ERR_ALIAS }, });
      }
  
      private final String groupbyInput;
  
      private final String expectedGroupingSets;
  
      private final IError expectedErrorCode;
  
      public SqlppGroupingSetsParserTest(String groupbyInput, Object expectedResult) {
          this.groupbyInput = groupbyInput;
          if (expectedResult instanceof IError) {
              this.expectedGroupingSets = null;
              this.expectedErrorCode = (IError) expectedResult;
          } else {
              this.expectedGroupingSets = (String) expectedResult;
              this.expectedErrorCode = null;
          }
      }
  
      @Test
      public void test() throws Exception {
          SqlppParserFactory parserFactory = new SqlppParserFactory(new NamespaceResolver(false));
          String groupbyClause = "GROUP BY " + groupbyInput;
          String query = "SELECT COUNT(*) FROM test " + groupbyClause + ";";
          // parse 2 queries so we can test calling rewrite() multiple times on the same instance
          IParser parser = parserFactory.createParser(query + query);
          List<Statement> statements;
          try {
              statements = parser.parse();
          } catch (CompilationException e) {
              if (expectedErrorCode == null) {
                  throw e;
              } else if (e.getMessage().contains(expectedErrorCode.errorCode())) {
                  return; // Found expected error code. SUCCESS
              } else {
                  Assert.fail(String.format("Unable to find expected error code %s in error message: %s",
                          expectedErrorCode, e.getMessage()));
                  throw new IllegalStateException();
              }
          }
          Assert.assertEquals(2, statements.size());
  
          for (Statement statement : statements) {
              String groupingSets = extractGroupingSets(statement);
              Assert.assertEquals(expectedGroupingSets, groupingSets);
          }
      }
  
      private String extractGroupingSets(Statement stmt) throws Exception {
          SqlppGroupingSetsRewriterTestVisitor visitor = new SqlppGroupingSetsRewriterTestVisitor();
          stmt.accept(visitor, null);
          return visitor.printGroupingSets();
      }
  
      private static class SqlppGroupingSetsRewriterTestVisitor extends AbstractSqlppSimpleExpressionVisitor {
  
          private final List<List<GbyVariableExpressionPair>> groupingSets = new ArrayList<>();
  
          private final StringWriter stringWriter = new StringWriter();
  
          private final PrintWriter printWriter = new PrintWriter(stringWriter);
  
          private final SqlppFormatPrintVisitor printVisitor = new SqlppFormatPrintVisitor(printWriter);
  
          @Override
          public Expression visit(GroupbyClause gc, ILangExpression arg) throws CompilationException {
              groupingSets.addAll(gc.getGbyPairList());
              return super.visit(gc, arg);
          }
  
          public String printGroupingSets() throws CompilationException {
              StringBuffer sb = stringWriter.getBuffer();
              sb.delete(0, sb.length());
  
              for (List<GbyVariableExpressionPair> groupingSet : groupingSets) {
                  printWriter.append('(');
                  String sep = "";
                  for (GbyVariableExpressionPair pair : groupingSet) {
                      printWriter.append(sep);
                      sep = ",";
                      pair.getExpr().accept(printVisitor, 0);
                      if (pair.getVar() != null) {
                          String ident = SqlppVariableUtil.toUserDefinedName(pair.getVar().getVar().getValue());
                          printWriter.append(" as ").append(ident);
                      }
                  }
                  printWriter.append(')');
              }
              printWriter.flush();
              return stringWriter.toString();
          }
      }
  
      private static String generateSimpleGroupingSet(String prefix, int n) {
          StringBuilder sb = new StringBuilder();
          for (int i = 0; i < n; i++) {
              if (i > 0) {
                  sb.append(',');
              }
              sb.append(prefix).append(i);
          }
          return sb.toString();
      }
  }
  