| /** |
| * 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.pinot.integration.tests; |
| |
| import com.fasterxml.jackson.databind.node.ObjectNode; |
| import java.io.BufferedWriter; |
| import java.io.File; |
| import java.io.FileWriter; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.Collections; |
| import java.util.HashMap; |
| import java.util.HashSet; |
| import java.util.Iterator; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Random; |
| import java.util.Set; |
| import org.apache.avro.Schema; |
| import org.apache.avro.file.DataFileReader; |
| import org.apache.avro.generic.GenericData; |
| import org.apache.avro.generic.GenericDatumReader; |
| import org.apache.avro.generic.GenericRecord; |
| import org.apache.commons.lang.StringUtils; |
| import org.apache.pinot.spi.utils.JsonUtils; |
| |
| |
| /** |
| * The <code>QueryGenerator</code> class is used to generate random equivalent Pinot and H2 query pairs based on Avro |
| * files. |
| * <ul> |
| * <li> |
| * Supports COMPARISON, IN and BETWEEN predicate for both single-value and multi-value columns. |
| * <ul> |
| * <li>For multi-value columns, does not support NOT EQUAL and NOT IN.</li> |
| * <li>The reason for this restriction is that number of elements in multi-value columns is not fixed.</li> |
| * </ul> |
| * </li> |
| * <li>Supports single-value data type: BOOLEAN, INT, LONG, FLOAT, DOUBLE, STRING.</li> |
| * <li>Supports multi-value data type: INT, LONG, FLOAT, DOUBLE, STRING.</li> |
| * <li> |
| * Supports aggregation function: SUM, MIN, MAX, AVG, COUNT, DISTINCTCOUNT. |
| * <ul> |
| * <li>SUM, MIN, MAX, AVG can only work on numeric single-value columns.</li> |
| * <li>COUNT, DISTINCTCOUNT can work on any single-value columns.</li> |
| * </ul> |
| * </li> |
| * </ul> |
| */ |
| public class QueryGenerator { |
| // Configurable variables. |
| private static final int MAX_NUM_SELECTION_COLUMNS = 3; |
| private static final int MAX_NUM_AGGREGATION_COLUMNS = 3; |
| private static final int MAX_NUM_ORDER_BY_COLUMNS = 3; |
| private static final int MAX_NUM_GROUP_BY_COLUMNS = 3; |
| private static final int MAX_NUM_PREDICATES = 3; |
| private static final int MAX_NUM_IN_CLAUSE_VALUES = 5; |
| private static final int MAX_RESULT_LIMIT = 30; |
| private static final int MAX_INT_FOR_HAVING_CLAUSE_PREDICATE = 1000; |
| private static final float MAX_FLOAT_FOR_HAVING_CLAUSE_PREDICATE = 10000; |
| private static final List<String> BOOLEAN_OPERATORS = Arrays.asList("OR", "AND"); |
| private static final List<String> COMPARISON_OPERATORS = Arrays.asList("=", "<>", "<", ">", "<=", ">="); |
| |
| private static final List<String> AGGREGATION_FUNCTIONS = |
| Arrays.asList("SUM", "MIN", "MAX", "AVG", "COUNT", "DISTINCTCOUNT"); |
| private static final List<String> MULTI_STAGE_AGGREGATION_FUNCTIONS = |
| Arrays.asList("SUM", "MIN", "MAX", "AVG", "COUNT"); |
| private static final Random RANDOM = new Random(); |
| |
| private final Map<String, Set<String>> _columnToValueSet = new HashMap<>(); |
| private final Map<String, List<String>> _columnToValueList = new HashMap<>(); |
| private final List<String> _columnNames = new ArrayList<>(); |
| private final List<String> _singleValueColumnNames = new ArrayList<>(); |
| private final List<String> _singleValueNumericalColumnNames = new ArrayList<>(); |
| private final Map<String, Integer> _multiValueColumnMaxNumElements = new HashMap<>(); |
| |
| private final List<QueryGenerationStrategy> _queryGenerationStrategies = |
| Arrays.asList(new SelectionQueryGenerationStrategy(), new AggregationQueryGenerationStrategy()); |
| private final List<PredicateGenerator> _singleValuePredicateGenerators = |
| Arrays.asList(new SingleValueComparisonPredicateGenerator(), new SingleValueInPredicateGenerator(), |
| new SingleValueBetweenPredicateGenerator(), new SingleValueRegexPredicateGenerator()); |
| private final List<PredicateGenerator> _multistageSingleValuePredicateGenerators = |
| Arrays.asList(new SingleValueComparisonPredicateGenerator(), new SingleValueInPredicateGenerator(), |
| new SingleValueBetweenPredicateGenerator()); |
| private final List<PredicateGenerator> _multiValuePredicateGenerators = |
| Arrays.asList(new MultiValueComparisonPredicateGenerator(), new MultiValueInPredicateGenerator(), |
| new MultiValueBetweenPredicateGenerator()); |
| |
| private final String _pinotTableName; |
| private final String _h2TableName; |
| private boolean _skipMultiValuePredicates = false; |
| private boolean _useMultistageEngine = false; |
| |
| /** |
| * Constructor for <code>QueryGenerator</code>. |
| * |
| * @param avroFiles list of Avro files. |
| * @param pinotTableName Pinot table name. |
| * @param h2TableName H2 table name. |
| */ |
| public QueryGenerator(List<File> avroFiles, String pinotTableName, String h2TableName) { |
| _pinotTableName = pinotTableName; |
| _h2TableName = h2TableName; |
| |
| // Read Avro schema and initialize storage. |
| File schemaAvroFile = avroFiles.get(0); |
| GenericDatumReader<GenericRecord> datumReader = new GenericDatumReader<>(); |
| try (DataFileReader<GenericRecord> fileReader = new DataFileReader<>(schemaAvroFile, datumReader)) { |
| Schema schema = fileReader.getSchema(); |
| for (Schema.Field field : schema.getFields()) { |
| String fieldName = field.name(); |
| Schema fieldSchema = field.schema(); |
| Schema.Type fieldType = fieldSchema.getType(); |
| |
| switch (fieldType) { |
| case UNION: |
| _columnNames.add(fieldName); |
| _columnToValueSet.put(fieldName, new HashSet<>()); |
| Schema.Type type = fieldSchema.getTypes().get(0).getType(); |
| if (type == Schema.Type.ARRAY) { |
| _multiValueColumnMaxNumElements.put(fieldName, 0); |
| } else { |
| _singleValueColumnNames.add(fieldName); |
| if (type != Schema.Type.STRING && type != Schema.Type.BOOLEAN) { |
| _singleValueNumericalColumnNames.add(fieldName); |
| } |
| } |
| break; |
| case ARRAY: |
| _columnNames.add(fieldName); |
| _columnToValueSet.put(fieldName, new HashSet<>()); |
| _multiValueColumnMaxNumElements.put(fieldName, 0); |
| break; |
| case INT: |
| case LONG: |
| case FLOAT: |
| case DOUBLE: |
| _columnNames.add(fieldName); |
| _columnToValueSet.put(fieldName, new HashSet<>()); |
| _singleValueColumnNames.add(fieldName); |
| _singleValueNumericalColumnNames.add(fieldName); |
| break; |
| case BOOLEAN: |
| case STRING: |
| _columnNames.add(fieldName); |
| _columnToValueSet.put(fieldName, new HashSet<>()); |
| _singleValueColumnNames.add(fieldName); |
| break; |
| default: |
| break; |
| } |
| } |
| } catch (Exception e) { |
| throw new RuntimeException(e); |
| } |
| |
| // Load Avro data into storage. |
| for (File avroFile : avroFiles) { |
| addAvroData(avroFile); |
| } |
| |
| // Ignore multi-value columns with too many elements. |
| prepareToGenerateQueries(); |
| } |
| |
| /** |
| * Helper method to store an Avro value into the valid SQL String value set. |
| * |
| * @param valueSet value set. |
| * @param avroValue Avro value. |
| */ |
| private static void storeAvroValueIntoValueSet(Set<String> valueSet, Object avroValue) { |
| if (avroValue instanceof Number) { |
| // For Number object, store raw value. |
| valueSet.add(avroValue.toString()); |
| } else { |
| // For non-Number object, escape single quote. |
| valueSet.add("'" + avroValue.toString().replace("'", "''") + "'"); |
| } |
| } |
| |
| /** |
| * Helper method to join several {@link String} elements with ' '. |
| * |
| * @param elements elements to be joined. |
| * @return joined result. |
| */ |
| private static String joinWithSpaces(String... elements) { |
| StringBuilder stringBuilder = new StringBuilder(); |
| for (String element : elements) { |
| if (!element.isEmpty()) { |
| stringBuilder.append(element).append(' '); |
| } |
| } |
| return stringBuilder.substring(0, stringBuilder.length() - 1); |
| } |
| |
| /** |
| * Sample main class for the query generator. |
| * |
| * @param args arguments. |
| */ |
| public static void main(String[] args) |
| throws Exception { |
| File avroFile = new File("pinot-integration-tests/src/test/resources/On_Time_On_Time_Performance_2014_1.avro"); |
| QueryGenerator queryGenerator = new QueryGenerator(Collections.singletonList(avroFile), "mytable", "mytable"); |
| File outputFile = new File( |
| "pinot-integration-tests/src/test/resources/On_Time_On_Time_Performance_2014_100k_subset.test_queries_10K.sql"); |
| try (BufferedWriter writer = new BufferedWriter(new FileWriter(outputFile))) { |
| for (int i = 0; i < 10000; i++) { |
| Query query = queryGenerator.generateQuery(); |
| ObjectNode queryJson = JsonUtils.newObjectNode(); |
| queryJson.put("sql", query.generatePinotQuery()); |
| queryJson.set("hsqls", JsonUtils.objectToJsonNode(query.generateH2Query())); |
| writer.write(queryJson.toString()); |
| writer.newLine(); |
| } |
| } |
| } |
| |
| /** |
| * Helper method to read in an Avro file and add data to the storage. |
| * |
| * @param avroFile Avro file. |
| */ |
| private void addAvroData(File avroFile) { |
| // Read in records and update the values stored. |
| GenericDatumReader<GenericRecord> datumReader = new GenericDatumReader<>(); |
| try (DataFileReader<GenericRecord> fileReader = new DataFileReader<>(avroFile, datumReader)) { |
| for (GenericRecord genericRecord : fileReader) { |
| for (String columnName : _columnNames) { |
| Set<String> values = _columnToValueSet.get(columnName); |
| |
| // Turn the Avro value into a valid SQL String token. |
| Object avroValue = genericRecord.get(columnName); |
| if (avroValue != null) { |
| Integer storedMaxNumElements = _multiValueColumnMaxNumElements.get(columnName); |
| if (storedMaxNumElements != null) { |
| // Multi-value column |
| GenericData.Array array = (GenericData.Array) avroValue; |
| int numElements = array.size(); |
| if (storedMaxNumElements < numElements) { |
| _multiValueColumnMaxNumElements.put(columnName, numElements); |
| } |
| for (Object element : array) { |
| storeAvroValueIntoValueSet(values, element); |
| } |
| } else { |
| // Single-value column |
| storeAvroValueIntoValueSet(values, avroValue); |
| } |
| } |
| } |
| } |
| } catch (Exception e) { |
| throw new RuntimeException(e); |
| } |
| } |
| |
| /** |
| * Helper method to finish initialization of the <code>QueryGenerator</code>, removing multi-value columns with too |
| * many elements and dumping storage into the final map from column name to list of column values. |
| * <p>Called after all Avro data loaded. |
| */ |
| private void prepareToGenerateQueries() { |
| Iterator<String> columnNameIterator = _columnNames.iterator(); |
| while (columnNameIterator.hasNext()) { |
| String columnName = columnNameIterator.next(); |
| |
| // Remove multi-value columns with more than MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE elements. |
| Integer maxNumElements = _multiValueColumnMaxNumElements.get(columnName); |
| if (maxNumElements != null |
| && maxNumElements > ClusterIntegrationTestUtils.MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE) { |
| columnNameIterator.remove(); |
| _multiValueColumnMaxNumElements.remove(columnName); |
| } else { |
| _columnToValueList.put(columnName, new ArrayList<>(_columnToValueSet.get(columnName))); |
| } |
| } |
| |
| // Free the other copy of the data. |
| _columnToValueSet.clear(); |
| } |
| |
| /** |
| * Set whether to skip predicates on multi-value columns. |
| * |
| * @param skipMultiValuePredicates whether to skip predicates on multi-value columns. |
| */ |
| public void setSkipMultiValuePredicates(boolean skipMultiValuePredicates) { |
| _skipMultiValuePredicates = skipMultiValuePredicates; |
| } |
| |
| public void setUseMultistageEngine(boolean useMultistageEngine) { |
| _useMultistageEngine = useMultistageEngine; |
| } |
| |
| /** |
| * Helper method to pick a random value from the values list. |
| * |
| * @param list values list. |
| * @param <T> type of the value. |
| * @return randomly picked value. |
| */ |
| private <T> T pickRandom(List<T> list) { |
| return list.get(RANDOM.nextInt(list.size())); |
| } |
| |
| private List<PredicateGenerator> getSingleValuePredicateGenerators() { |
| return _useMultistageEngine ? _multistageSingleValuePredicateGenerators : _singleValuePredicateGenerators; |
| } |
| |
| /** |
| * Generate one selection or aggregation query. |
| * |
| * @return generated query. |
| */ |
| public Query generateQuery() { |
| return pickRandom(_queryGenerationStrategies).generateQuery(); |
| } |
| |
| /** |
| * Helper method to generate a predicate query fragment. |
| * |
| * @return generated predicate query fragment. |
| */ |
| private PredicateQueryFragment generatePredicate() { |
| // Generate at most MAX_NUM_PREDICATES predicates. |
| int predicateCount = RANDOM.nextInt(MAX_NUM_PREDICATES + 1); |
| |
| List<QueryFragment> predicates = new ArrayList<>(predicateCount); |
| while (predicates.size() < predicateCount) { |
| String columnName = pickRandom(_columnNames); |
| if (!_columnToValueList.get(columnName).isEmpty()) { |
| if (!_multiValueColumnMaxNumElements.containsKey(columnName)) { |
| // Single-value column. |
| predicates.add(pickRandom(getSingleValuePredicateGenerators()).generatePredicate(columnName)); |
| } else if (!_skipMultiValuePredicates) { |
| // Multi-value column. |
| predicates.add(pickRandom(_multiValuePredicateGenerators).generatePredicate(columnName)); |
| } |
| } |
| } |
| |
| if (predicateCount < 2) { |
| // No need to join. |
| return new PredicateQueryFragment(predicates, Collections.emptyList()); |
| } else { |
| // Join predicates with ANDs and ORs. |
| List<QueryFragment> operators = new ArrayList<>(predicateCount - 1); |
| for (int i = 1; i < predicateCount; i++) { |
| operators.add(new StringQueryFragment(pickRandom(BOOLEAN_OPERATORS))); |
| } |
| return new PredicateQueryFragment(predicates, operators); |
| } |
| } |
| |
| public interface Query { |
| |
| String generatePinotQuery(); |
| |
| String generateH2Query(); |
| } |
| |
| private interface QueryFragment { |
| |
| String generatePinotQuery(); |
| |
| String generateH2Query(); |
| } |
| |
| /** |
| * Query generation strategy interface with capability of generating query using specific strategy. |
| */ |
| private interface QueryGenerationStrategy { |
| |
| /** |
| * Generate a query using specific strategy. |
| * |
| * @return generated query. |
| */ |
| Query generateQuery(); |
| } |
| |
| /** |
| * Predicate generator interface with capability of generating a predicate query fragment on a column. |
| */ |
| private interface PredicateGenerator { |
| |
| /** |
| * Generate a predicate query fragment on a column. |
| * |
| * @param columnName column name. |
| * @return generated predicate query fragment. |
| */ |
| QueryFragment generatePredicate(String columnName); |
| } |
| |
| /** |
| * Selection query. |
| */ |
| private class SelectionQuery implements Query { |
| final List<String> _projectionColumns; |
| final PredicateQueryFragment _predicate; |
| final OrderByQueryFragment _orderBy; |
| final LimitQueryFragment _limit; |
| |
| /** |
| * Constructor for <code>SelectionQuery</code>. |
| * |
| * @param projectionColumns projection columns. |
| * @param orderBy order by fragment. |
| * @param predicate predicate fragment. |
| * @param limit limit fragment. |
| */ |
| public SelectionQuery(List<String> projectionColumns, PredicateQueryFragment predicate, |
| OrderByQueryFragment orderBy, LimitQueryFragment limit) { |
| _projectionColumns = projectionColumns; |
| _orderBy = orderBy; |
| _predicate = predicate; |
| _limit = limit; |
| } |
| |
| @Override |
| public String generatePinotQuery() { |
| return joinWithSpaces("SELECT", StringUtils.join(_projectionColumns, ", "), "FROM", _pinotTableName, |
| _predicate.generatePinotQuery(), _orderBy.generatePinotQuery(), _limit.generatePinotQuery()); |
| } |
| |
| @Override |
| public String generateH2Query() { |
| List<String> h2ProjectionColumns = new ArrayList<>(); |
| for (String projectionColumn : _projectionColumns) { |
| if (_multiValueColumnMaxNumElements.containsKey(projectionColumn)) { |
| // Multi-value column. |
| for (int i = 0; i < ClusterIntegrationTestUtils.MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE; i++) { |
| h2ProjectionColumns.add(projectionColumn + "__MV" + i); |
| } |
| } else { |
| // Single-value column. |
| h2ProjectionColumns.add(projectionColumn); |
| } |
| } |
| return joinWithSpaces("SELECT", StringUtils.join(h2ProjectionColumns, ", "), "FROM", _h2TableName, |
| _predicate.generateH2Query(), _orderBy.generateH2Query(), _limit.generateH2Query()); |
| } |
| } |
| |
| /** |
| * Aggregation query. |
| */ |
| private class AggregationQuery implements Query { |
| final List<String> _aggregateColumnsAndFunctions; |
| final PredicateQueryFragment _predicate; |
| final HavingQueryFragment _havingPredicate; |
| final Set<String> _groupColumns; |
| final LimitQueryFragment _limit; |
| |
| /** |
| * Constructor for <code>AggregationQuery</code>. |
| * |
| * @param aggregateColumnsAndFunctions aggregation functions. |
| * @param predicate predicate fragment. |
| * @param groupColumns group-by columns. |
| * @param havingPredicate having predicate fragment. |
| * @param limit limit fragment. |
| */ |
| public AggregationQuery(List<String> aggregateColumnsAndFunctions, PredicateQueryFragment predicate, |
| Set<String> groupColumns, HavingQueryFragment havingPredicate, LimitQueryFragment limit) { |
| _aggregateColumnsAndFunctions = aggregateColumnsAndFunctions; |
| _predicate = predicate; |
| _groupColumns = groupColumns; |
| _havingPredicate = havingPredicate; |
| _limit = limit; |
| } |
| |
| @Override |
| public String generatePinotQuery() { |
| if (_groupColumns.isEmpty()) { |
| return joinWithSpaces("SELECT", StringUtils.join(_aggregateColumnsAndFunctions, ", "), "FROM", _pinotTableName, |
| _predicate.generatePinotQuery()); |
| } else { |
| return joinWithSpaces("SELECT", StringUtils.join(_aggregateColumnsAndFunctions, ", "), "FROM", _pinotTableName, |
| _predicate.generatePinotQuery(), "GROUP BY", StringUtils.join(_groupColumns, ", "), |
| _havingPredicate.generatePinotQuery(), _limit.generatePinotQuery()); |
| } |
| } |
| |
| @Override |
| public String generateH2Query() { |
| List<String> h2AggregateColumnAndFunctions = new ArrayList<>(); |
| for (String aggregateColumnAndFunction : _aggregateColumnsAndFunctions) { |
| String h2AggregateColumnAndFunction; |
| // Make 'AVG' and |
| if (aggregateColumnAndFunction.startsWith("DISTINCTCOUNT(")) { |
| // make 'DISTINCTCOUNT(..)' compatible with H2 SQL query using 'COUNT(DISTINCT(..)' |
| h2AggregateColumnAndFunction = aggregateColumnAndFunction.replace("DISTINCTCOUNT(", "COUNT(DISTINCT "); |
| } else if (AGGREGATION_FUNCTIONS.contains(aggregateColumnAndFunction.substring(0, 3))) { |
| // make AGG functions (SUM, MIN, MAX, AVG) compatible with H2 SQL query. |
| // this is because Pinot queries casts all to double before doing aggregation |
| String aggFunctionName = aggregateColumnAndFunction.substring(0, 3); |
| h2AggregateColumnAndFunction = aggregateColumnAndFunction |
| .replace(aggFunctionName + "(", aggFunctionName + "(CAST(") |
| .replace(")", " AS DOUBLE))"); |
| } else { |
| h2AggregateColumnAndFunction = aggregateColumnAndFunction; |
| } |
| h2AggregateColumnAndFunctions.add(h2AggregateColumnAndFunction); |
| } |
| if (_groupColumns.isEmpty()) { |
| return joinWithSpaces("SELECT", StringUtils.join(h2AggregateColumnAndFunctions, ", "), "FROM", _h2TableName, |
| _predicate.generateH2Query()); |
| } else { |
| return joinWithSpaces("SELECT", StringUtils.join(h2AggregateColumnAndFunctions, ", "), "FROM", _h2TableName, |
| _predicate.generateH2Query(), "GROUP BY", StringUtils.join(_groupColumns, ", "), |
| _havingPredicate.generateH2Query(), _limit.generateH2Query()); |
| } |
| } |
| } |
| |
| /** |
| * Most basic query fragment. |
| */ |
| private static class StringQueryFragment implements QueryFragment { |
| final String _pinotQuery; |
| final String _h2Query; |
| |
| /** |
| * Constructor with same Pinot and H2 query fragment. |
| */ |
| StringQueryFragment(String query) { |
| _pinotQuery = query; |
| _h2Query = query; |
| } |
| |
| /** |
| * Constructor for <code>StringQueryFragment</code> with different Pinot and H2 query fragment. |
| */ |
| StringQueryFragment(String pinotQuery, String h2Query) { |
| _pinotQuery = pinotQuery; |
| _h2Query = h2Query; |
| } |
| |
| @Override |
| public String generatePinotQuery() { |
| return _pinotQuery; |
| } |
| |
| @Override |
| public String generateH2Query() { |
| return _h2Query; |
| } |
| } |
| |
| /** |
| * Limit query fragment for selection queries. |
| * <ul> |
| * <li>SELECT ... FROM ... WHERE ... 'LIMIT ...'</li> |
| * </ul> |
| */ |
| private static class LimitQueryFragment extends StringQueryFragment { |
| LimitQueryFragment(int limit) { |
| // When limit is MAX_RESULT_LIMIT, construct query without LIMIT. |
| super(limit == MAX_RESULT_LIMIT ? "" : "LIMIT " + limit, |
| "LIMIT " + ClusterIntegrationTestUtils.MAX_NUM_ROWS_TO_COMPARE); |
| } |
| } |
| |
| /** |
| * Order by query fragment for aggregation queries. |
| * <ul> |
| * <li>SELECT ... FROM ... WHERE ... 'ORDER BY ...'</li> |
| * </ul> |
| */ |
| private static class OrderByQueryFragment extends StringQueryFragment { |
| OrderByQueryFragment(Set<String> columns) { |
| super(columns.isEmpty() ? "" : "ORDER BY " + StringUtils.join(columns, ", ")); |
| } |
| } |
| |
| /** |
| * Predicate query fragment. |
| * <ul> |
| * <li>SELECT ... FROM ... 'WHERE ...'</li> |
| * </ul> |
| */ |
| private static class PredicateQueryFragment implements QueryFragment { |
| final List<QueryFragment> _predicates; |
| final List<QueryFragment> _operators; |
| |
| /** |
| * Constructor for <code>PredicateQueryFragment</code>. |
| * |
| * @param predicates predicates. |
| * @param operators operators between predicates. |
| */ |
| PredicateQueryFragment(List<QueryFragment> predicates, List<QueryFragment> operators) { |
| _predicates = predicates; |
| _operators = operators; |
| } |
| |
| @Override |
| public String generatePinotQuery() { |
| if (_predicates.isEmpty()) { |
| return ""; |
| } else { |
| StringBuilder pinotQuery = new StringBuilder("WHERE "); |
| |
| // One less than the number of predicates. |
| int operatorCount = _operators.size(); |
| for (int i = 0; i < operatorCount; i++) { |
| pinotQuery.append(_predicates.get(i).generatePinotQuery()).append(' ') |
| .append(_operators.get(i).generatePinotQuery()).append(' '); |
| } |
| pinotQuery.append(_predicates.get(operatorCount).generatePinotQuery()); |
| |
| return pinotQuery.toString(); |
| } |
| } |
| |
| @Override |
| public String generateH2Query() { |
| if (_predicates.isEmpty()) { |
| return ""; |
| } else { |
| StringBuilder h2Query = new StringBuilder("WHERE "); |
| |
| // One less than the number of predicates. |
| int operatorCount = _operators.size(); |
| for (int i = 0; i < operatorCount; i++) { |
| h2Query.append(_predicates.get(i).generateH2Query()).append(' ').append(_operators.get(i).generateH2Query()) |
| .append(' '); |
| } |
| h2Query.append(_predicates.get(operatorCount).generateH2Query()); |
| |
| return h2Query.toString(); |
| } |
| } |
| } |
| |
| /** |
| * Having query fragment. |
| * <ul> |
| * <li>SELECT ... FROM ... WHERE ... GROUP BY ... 'HAVING ...'</li> |
| * </ul> |
| */ |
| private static class HavingQueryFragment implements QueryFragment { |
| final List<String> _havingClauseAggregationFunctions; |
| final List<String> _havingClauseOperatorsAndValues; |
| final List<String> _havingClauseBooleanOperators; |
| |
| HavingQueryFragment(List<String> havingClauseAggregationFunctions, List<String> havingClauseOperatorsAndValues, |
| List<String> havingClauseBooleanOperators) { |
| _havingClauseAggregationFunctions = havingClauseAggregationFunctions; |
| _havingClauseOperatorsAndValues = havingClauseOperatorsAndValues; |
| _havingClauseBooleanOperators = havingClauseBooleanOperators; |
| } |
| |
| @Override |
| public String generatePinotQuery() { |
| String pinotQuery = ""; |
| int aggregationFunctionCount = _havingClauseAggregationFunctions.size(); |
| if (aggregationFunctionCount > 0) { |
| pinotQuery += "HAVING"; |
| pinotQuery = joinWithSpaces(pinotQuery, _havingClauseAggregationFunctions.get(0)); |
| pinotQuery = joinWithSpaces(pinotQuery, _havingClauseOperatorsAndValues.get(0)); |
| for (int i = 1; i < aggregationFunctionCount; i++) { |
| pinotQuery = joinWithSpaces(pinotQuery, _havingClauseBooleanOperators.get(i - 1)); |
| pinotQuery = joinWithSpaces(pinotQuery, _havingClauseAggregationFunctions.get(i)); |
| pinotQuery = joinWithSpaces(pinotQuery, _havingClauseOperatorsAndValues.get(i)); |
| } |
| } |
| return pinotQuery; |
| } |
| |
| @Override |
| public String generateH2Query() { |
| String h2Query = ""; |
| int aggregationFunctionCount = _havingClauseAggregationFunctions.size(); |
| if (aggregationFunctionCount > 0) { |
| h2Query += "HAVING"; |
| String aggregationFunction = _havingClauseAggregationFunctions.get(0); |
| if (aggregationFunction.startsWith("AVG(")) { |
| aggregationFunction = aggregationFunction.replace("AVG(", "AVG(CAST(").replace(")", " AS DOUBLE))"); |
| } else if (aggregationFunction.startsWith("DISTINCTCOUNT(")) { |
| aggregationFunction = aggregationFunction.replace("DISTINCTCOUNT(", "COUNT(DISTINCT "); |
| } |
| h2Query = joinWithSpaces(h2Query, aggregationFunction); |
| h2Query = joinWithSpaces(h2Query, _havingClauseOperatorsAndValues.get(0)); |
| for (int i = 1; i < aggregationFunctionCount; i++) { |
| aggregationFunction = _havingClauseAggregationFunctions.get(i); |
| if (aggregationFunction.startsWith("AVG(")) { |
| aggregationFunction = aggregationFunction.replace("AVG(", "AVG(CAST(").replace(")", " AS DOUBLE))"); |
| } else if (aggregationFunction.startsWith("DISTINCTCOUNT(")) { |
| aggregationFunction = aggregationFunction.replace("DISTINCTCOUNT(", "COUNT(DISTINCT "); |
| } |
| h2Query = joinWithSpaces(h2Query, _havingClauseBooleanOperators.get(i - 1)); |
| h2Query = joinWithSpaces(h2Query, aggregationFunction); |
| h2Query = joinWithSpaces(h2Query, _havingClauseOperatorsAndValues.get(i)); |
| } |
| } |
| return h2Query; |
| } |
| } |
| |
| /** |
| * Strategy to generate selection queries. |
| * <ul> |
| * <li>SELECT a, b FROM table WHERE a = 'foo' AND b = 'bar' ORDER BY c LIMIT 10</li> |
| * </ul> |
| */ |
| private class SelectionQueryGenerationStrategy implements QueryGenerationStrategy { |
| |
| @Override |
| public Query generateQuery() { |
| // Select at most MAX_NUM_SELECTION_COLUMNS columns. |
| int projectionColumnCount = Math.min(RANDOM.nextInt(MAX_NUM_SELECTION_COLUMNS) + 1, _columnNames.size()); |
| if (_useMultistageEngine) { |
| projectionColumnCount = Math.min(projectionColumnCount, _singleValueColumnNames.size()); |
| } |
| Set<String> projectionColumns = new HashSet<>(); |
| while (projectionColumns.size() < projectionColumnCount) { |
| projectionColumns.add(pickRandom(_useMultistageEngine ? _singleValueColumnNames : _columnNames)); |
| } |
| |
| // Select at most MAX_NUM_ORDER_BY_COLUMNS columns for ORDER BY clause. |
| int orderByColumnCount = Math.min(RANDOM.nextInt(MAX_NUM_ORDER_BY_COLUMNS + 1), _singleValueColumnNames.size()); |
| Set<String> orderByColumns = new HashSet<>(); |
| while (orderByColumns.size() < orderByColumnCount) { |
| orderByColumns.add(pickRandom(_singleValueColumnNames)); |
| } |
| |
| // Generate a predicate. |
| PredicateQueryFragment predicate = generatePredicate(); |
| |
| // Generate a result limit of at most MAX_RESULT_LIMIT columns for ORDER BY clause. |
| int resultLimit = RANDOM.nextInt(MAX_RESULT_LIMIT + 1); |
| LimitQueryFragment limit = new LimitQueryFragment(resultLimit); |
| |
| return new SelectionQuery(new ArrayList<>(projectionColumns), predicate, new OrderByQueryFragment(orderByColumns), |
| limit); |
| } |
| } |
| |
| /** |
| * Strategy to generate aggregation queries. |
| * <ul> |
| * <li>SELECT SUM(a), MAX(b) FROM table WHERE a = 'foo' AND b = 'bar' GROUP BY c TOP 10</li> |
| * </ul> |
| */ |
| private class AggregationQueryGenerationStrategy implements QueryGenerationStrategy { |
| |
| @Override |
| public Query generateQuery() { |
| // Generate at most MAX_NUM_AGGREGATION_COLUMNS columns on which to aggregate, map 0 to 'COUNT(*)'. |
| int aggregationColumnCount = RANDOM.nextInt(MAX_NUM_AGGREGATION_COLUMNS + 1); |
| Set<String> aggregationColumnsAndFunctions = new HashSet<>(); |
| if (aggregationColumnCount == 0) { |
| aggregationColumnsAndFunctions.add("COUNT(*)"); |
| } else { |
| while (aggregationColumnsAndFunctions.size() < aggregationColumnCount) { |
| aggregationColumnsAndFunctions.add(createRandomAggregationFunction()); |
| } |
| } |
| // Generate a predicate. |
| PredicateQueryFragment predicate = generatePredicate(); |
| // Generate at most MAX_NUM_GROUP_BY_COLUMNS columns on which to group. |
| int groupColumnCount = Math.min(RANDOM.nextInt(MAX_NUM_GROUP_BY_COLUMNS + 1), _singleValueColumnNames.size()); |
| Set<String> groupColumns = new HashSet<>(); |
| while (groupColumns.size() < groupColumnCount) { |
| groupColumns.add(pickRandom(_singleValueColumnNames)); |
| } |
| //Generate a HAVING predicate |
| ArrayList<String> arrayOfAggregationColumnsAndFunctions = new ArrayList<>(aggregationColumnsAndFunctions); |
| HavingQueryFragment havingPredicate = generateHavingPredicate(arrayOfAggregationColumnsAndFunctions); |
| // Generate a result limit of at most MAX_RESULT_LIMIT. |
| LimitQueryFragment limit = new LimitQueryFragment(RANDOM.nextInt(MAX_RESULT_LIMIT + 1)); |
| return new AggregationQuery(arrayOfAggregationColumnsAndFunctions, predicate, groupColumns, havingPredicate, |
| limit); |
| } |
| |
| /** |
| * Helper method to generate a having predicate query fragment. |
| * |
| * @return generated predicate query fragment. |
| */ |
| private HavingQueryFragment generateHavingPredicate(ArrayList<String> arrayOfAggregationColumnsAndFunctions) { |
| //Generate a HAVING clause for group by query |
| List<String> havingClauseAggregationFunctions = new ArrayList<>(); |
| List<String> havingClauseOperatorsAndValues = new ArrayList<>(); |
| List<String> havingClauseBooleanOperators = new ArrayList<>(); |
| createHavingClause(arrayOfAggregationColumnsAndFunctions, havingClauseAggregationFunctions, |
| havingClauseOperatorsAndValues, havingClauseBooleanOperators); |
| return new HavingQueryFragment(havingClauseAggregationFunctions, havingClauseOperatorsAndValues, |
| havingClauseBooleanOperators); |
| } |
| |
| private void createHavingClause(ArrayList<String> arrayOfAggregationColumnsAndFunctions, |
| List<String> havingClauseAggregationFunctions, List<String> havingClauseOperatorsAndValues, |
| List<String> havingClauseBooleanOperators) { |
| int numOfFunctionsInSelectList = arrayOfAggregationColumnsAndFunctions.size(); |
| int aggregationFunctionCount = RANDOM.nextInt(numOfFunctionsInSelectList + 1); |
| for (int i = 0; i < aggregationFunctionCount; i++) { |
| int aggregationFunctionIndex = RANDOM.nextInt(numOfFunctionsInSelectList); |
| String aggregationFunction = arrayOfAggregationColumnsAndFunctions.get(aggregationFunctionIndex); |
| havingClauseAggregationFunctions.add(aggregationFunction); |
| havingClauseOperatorsAndValues.add(createOperatorAndValueForAggregationFunction(aggregationFunction)); |
| } |
| int aggregationFunctionNotInSelectListCount = RANDOM.nextInt(MAX_NUM_AGGREGATION_COLUMNS + 1); |
| for (int i = 0; i < aggregationFunctionNotInSelectListCount; i++) { |
| String aggregationFunction = createRandomAggregationFunction(); |
| havingClauseAggregationFunctions.add(aggregationFunction); |
| havingClauseOperatorsAndValues.add(createOperatorAndValueForAggregationFunction(aggregationFunction)); |
| } |
| aggregationFunctionCount += aggregationFunctionNotInSelectListCount; |
| for (int i = 1; i < aggregationFunctionCount; i++) { |
| havingClauseBooleanOperators.add(pickRandom(BOOLEAN_OPERATORS)); |
| } |
| } |
| |
| private String createRandomAggregationFunction() { |
| String aggregationFunction = pickRandom(_useMultistageEngine ? MULTI_STAGE_AGGREGATION_FUNCTIONS |
| : AGGREGATION_FUNCTIONS); |
| String aggregationColumn; |
| switch (aggregationFunction) { |
| // "COUNT" and "DISTINCTCOUNT" support all single-value columns. |
| case "COUNT": |
| case "DISTINCTCOUNT": |
| aggregationColumn = pickRandom(_singleValueColumnNames); |
| break; |
| // Other functions only support single-value numeric columns. |
| default: |
| aggregationColumn = pickRandom(_singleValueNumericalColumnNames); |
| break; |
| } |
| return aggregationFunction + "(" + aggregationColumn + ")"; |
| } |
| |
| private String createOperatorAndValueForAggregationFunction(String aggregationFunction) { |
| boolean generateInt = aggregationFunction.startsWith("COUNT") || aggregationFunction.startsWith("DISTINCTCOUNT"); |
| boolean comparisonOperatorNotBetweenOrIN = RANDOM.nextBoolean(); |
| if (comparisonOperatorNotBetweenOrIN) { |
| return pickRandom(COMPARISON_OPERATORS) + " " + generateRandomValue(generateInt); |
| } else { |
| boolean isItBetween = RANDOM.nextBoolean(); |
| if (isItBetween) { |
| return "BETWEEN " + generateRandomValue(generateInt) + " AND " + generateRandomValue(generateInt); |
| } else { |
| int numValues = RANDOM.nextInt(MAX_NUM_IN_CLAUSE_VALUES) + 1; |
| Set<String> values = new HashSet<>(); |
| while (values.size() < numValues) { |
| values.add(generateRandomValue(generateInt)); |
| } |
| String valuesString = StringUtils.join(values, ", "); |
| boolean isItIn = RANDOM.nextBoolean(); |
| return (isItIn ? "" : "NOT ") + "IN (" + valuesString + ")"; |
| } |
| } |
| } |
| |
| private String generateRandomValue(boolean generateInt) { |
| return generateInt ? Integer.toString(RANDOM.nextInt(MAX_INT_FOR_HAVING_CLAUSE_PREDICATE) + 1) |
| : Float.toString(RANDOM.nextFloat() * MAX_FLOAT_FOR_HAVING_CLAUSE_PREDICATE); |
| } |
| } |
| |
| /** |
| * Generator for single-value column comparison predicate query fragment. |
| */ |
| private class SingleValueComparisonPredicateGenerator implements PredicateGenerator { |
| |
| @Override |
| public QueryFragment generatePredicate(String columnName) { |
| String columnValue = pickRandom(_columnToValueList.get(columnName)); |
| String comparisonOperator = pickRandom(COMPARISON_OPERATORS); |
| return new StringQueryFragment(joinWithSpaces(columnName, comparisonOperator, columnValue)); |
| } |
| } |
| |
| /** |
| * Generator for single-value column <code>IN</code> predicate query fragment. |
| */ |
| private class SingleValueInPredicateGenerator implements PredicateGenerator { |
| |
| @Override |
| public QueryFragment generatePredicate(String columnName) { |
| List<String> columnValues = _columnToValueList.get(columnName); |
| |
| int numValues = Math.min(RANDOM.nextInt(MAX_NUM_IN_CLAUSE_VALUES) + 1, columnValues.size()); |
| Set<String> values = new HashSet<>(); |
| while (values.size() < numValues) { |
| values.add(pickRandom(columnValues)); |
| } |
| String inValues = StringUtils.join(values, ", "); |
| |
| boolean notIn = RANDOM.nextBoolean(); |
| if (notIn) { |
| return new StringQueryFragment(columnName + " NOT IN (" + inValues + ")"); |
| } else { |
| return new StringQueryFragment(columnName + " IN (" + inValues + ")"); |
| } |
| } |
| } |
| |
| /** |
| * Generator for single-value column <code>BETWEEN</code> predicate query fragment. |
| */ |
| private class SingleValueBetweenPredicateGenerator implements PredicateGenerator { |
| |
| @Override |
| public QueryFragment generatePredicate(String columnName) { |
| List<String> columnValues = _columnToValueList.get(columnName); |
| String leftValue = pickRandom(columnValues); |
| String rightValue = pickRandom(columnValues); |
| return new StringQueryFragment(columnName + " BETWEEN " + leftValue + " AND " + rightValue); |
| } |
| } |
| |
| /** |
| * Generator for single-value column <code>REGEX</code> predicate query fragment. |
| */ |
| private class SingleValueRegexPredicateGenerator implements PredicateGenerator { |
| Random _random = new Random(); |
| |
| @Override |
| public QueryFragment generatePredicate(String columnName) { |
| List<String> columnValues = _columnToValueList.get(columnName); |
| String value = pickRandom(columnValues); |
| // do regex only for string type |
| if (value.startsWith("'") && value.endsWith("'")) { |
| // replace only one character for now with .* ignore the first and last character |
| int indexToReplaceWithRegex = 1 + _random.nextInt(value.length() - 2); |
| String regex = value.substring(1, indexToReplaceWithRegex) + ".*" + value.substring(indexToReplaceWithRegex + 1, |
| value.length() - 1); |
| return new StringQueryFragment(String.format(" REGEXP_LIKE(%s, '%s')", columnName, regex), |
| String.format(" REGEXP_LIKE(%s, '%s', 'i')", columnName, regex)); |
| } else { |
| String equalsPredicate = String.format("%s = %s", columnName, value); |
| return new StringQueryFragment(equalsPredicate); |
| } |
| } |
| } |
| |
| /** |
| * Generator for multi-value column comparison predicate query fragment. |
| * <p>DO NOT SUPPORT '<code>NOT EQUAL</code>'. |
| */ |
| private class MultiValueComparisonPredicateGenerator implements PredicateGenerator { |
| |
| @Override |
| public QueryFragment generatePredicate(String columnName) { |
| String columnValue = pickRandom(_columnToValueList.get(columnName)); |
| String comparisonOperator = pickRandom(COMPARISON_OPERATORS); |
| |
| // Not equal works differently on multi-value, so avoid '<>' comparison. |
| while (comparisonOperator.equals("<>")) { |
| comparisonOperator = pickRandom(COMPARISON_OPERATORS); |
| } |
| |
| List<String> h2ComparisonClauses = |
| new ArrayList<>(ClusterIntegrationTestUtils.MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE); |
| for (int i = 0; i < ClusterIntegrationTestUtils.MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE; i++) { |
| h2ComparisonClauses.add(joinWithSpaces(columnName + "__MV" + i, comparisonOperator, columnValue)); |
| } |
| |
| return new StringQueryFragment(joinWithSpaces(columnName, comparisonOperator, columnValue), |
| "(" + StringUtils.join(h2ComparisonClauses, " OR ") + ")"); |
| } |
| } |
| |
| /** |
| * Generator for multi-value column <code>IN</code> predicate query fragment. |
| * <p>DO NOT SUPPORT '<code>NOT IN</code>'. |
| */ |
| private class MultiValueInPredicateGenerator implements PredicateGenerator { |
| |
| @Override |
| public QueryFragment generatePredicate(String columnName) { |
| List<String> columnValues = _columnToValueList.get(columnName); |
| |
| int numValues = Math.min(RANDOM.nextInt(MAX_NUM_IN_CLAUSE_VALUES) + 1, columnValues.size()); |
| Set<String> values = new HashSet<>(); |
| while (values.size() < numValues) { |
| values.add(pickRandom(columnValues)); |
| } |
| String inValues = StringUtils.join(values, ", "); |
| |
| List<String> h2InClauses = |
| new ArrayList<>(ClusterIntegrationTestUtils.MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE); |
| for (int i = 0; i < ClusterIntegrationTestUtils.MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE; i++) { |
| h2InClauses.add(columnName + "__MV" + i + " IN (" + inValues + ")"); |
| } |
| |
| return new StringQueryFragment(columnName + " IN (" + inValues + ")", |
| "(" + StringUtils.join(h2InClauses, " OR ") + ")"); |
| } |
| } |
| |
| /** |
| * Generator for multi-value column <code>BETWEEN</code> predicate query fragment. |
| */ |
| private class MultiValueBetweenPredicateGenerator implements PredicateGenerator { |
| |
| @Override |
| public QueryFragment generatePredicate(String columnName) { |
| List<String> columnValues = _columnToValueList.get(columnName); |
| String leftValue = pickRandom(columnValues); |
| String rightValue = pickRandom(columnValues); |
| |
| List<String> h2ComparisonClauses = |
| new ArrayList<>(ClusterIntegrationTestUtils.MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE); |
| for (int i = 0; i < ClusterIntegrationTestUtils.MAX_NUM_ELEMENTS_IN_MULTI_VALUE_TO_COMPARE; i++) { |
| h2ComparisonClauses.add(columnName + "__MV" + i + " BETWEEN " + leftValue + " AND " + rightValue); |
| } |
| |
| return new StringQueryFragment(columnName + " BETWEEN " + leftValue + " AND " + rightValue, |
| "(" + StringUtils.join(h2ComparisonClauses, " OR ") + ")"); |
| } |
| } |
| } |
