Examples with BrokerRequest com.linkedin.pinot.common.request.BrokerRequest used on opensource projects

Example 61 with BrokerRequest

use of com.linkedin.pinot.common.request.BrokerRequest in project pinot by linkedin.

the class OptimizationFlagsTest method testForbidEnableAndDisableTogether.

@Test
public void testForbidEnableAndDisableTogether() {
    BrokerRequest brokerRequest = new BrokerRequest();
    brokerRequest.setDebugOptions(Collections.singletonMap("optimizationFlags", "-foo, +bar"));
    try {
        OptimizationFlags optimizationFlags = OptimizationFlags.getOptimizationFlags(brokerRequest);
        Assert.fail("Expected exception");
    } catch (Exception e) {
    // Expected
    }
}

Example 62 with BrokerRequest

use of com.linkedin.pinot.common.request.BrokerRequest in project pinot by linkedin.

the class OptimizationFlagsTest method testEnableOptimizations.

@Test
public void testEnableOptimizations() {
    BrokerRequest brokerRequest = new BrokerRequest();
    brokerRequest.setDebugOptions(Collections.singletonMap("optimizationFlags", "+foo, +bar"));
    OptimizationFlags optimizationFlags = OptimizationFlags.getOptimizationFlags(brokerRequest);
    Assert.assertTrue(optimizationFlags.isOptimizationEnabled("foo"));
    Assert.assertTrue(optimizationFlags.isOptimizationEnabled("bar"));
    Assert.assertFalse(optimizationFlags.isOptimizationEnabled("baz"));
}

Example 63 with BrokerRequest

use of com.linkedin.pinot.common.request.BrokerRequest in project pinot by linkedin.

the class TransformGroupByTest method executeGroupByQuery.

/**
   * Helper method that executes the group by query on the index and returns the group by result.
   *
   * @param query Query to execute
   * @return Group by result
   */
private AggregationGroupByResult executeGroupByQuery(IndexSegment indexSegment, String query) {
    Operator filterOperator = new MatchEntireSegmentOperator(indexSegment.getSegmentMetadata().getTotalDocs());
    final BReusableFilteredDocIdSetOperator docIdSetOperator = new BReusableFilteredDocIdSetOperator(filterOperator, indexSegment.getSegmentMetadata().getTotalDocs(), NUM_ROWS);
    final Map<String, BaseOperator> dataSourceMap = buildDataSourceMap(indexSegment.getSegmentMetadata().getSchema());
    final MProjectionOperator projectionOperator = new MProjectionOperator(dataSourceMap, docIdSetOperator);
    Pql2Compiler compiler = new Pql2Compiler();
    BrokerRequest brokerRequest = compiler.compileToBrokerRequest(query);
    List<AggregationInfo> aggregationsInfo = brokerRequest.getAggregationsInfo();
    int numAggFunctions = aggregationsInfo.size();
    AggregationFunctionContext[] aggrFuncContextArray = new AggregationFunctionContext[numAggFunctions];
    AggregationFunctionInitializer aggFuncInitializer = new AggregationFunctionInitializer(indexSegment.getSegmentMetadata());
    for (int i = 0; i < numAggFunctions; i++) {
        AggregationInfo aggregationInfo = aggregationsInfo.get(i);
        aggrFuncContextArray[i] = AggregationFunctionContext.instantiate(aggregationInfo);
        aggrFuncContextArray[i].getAggregationFunction().accept(aggFuncInitializer);
    }
    GroupBy groupBy = brokerRequest.getGroupBy();
    Set<String> expressions = new HashSet<>(groupBy.getExpressions());
    TransformExpressionOperator transformOperator = new TransformExpressionOperator(projectionOperator, TransformPlanNode.buildTransformExpressionTrees(expressions));
    AggregationGroupByOperator groupByOperator = new AggregationGroupByOperator(aggrFuncContextArray, groupBy, Integer.MAX_VALUE, transformOperator, NUM_ROWS);
    IntermediateResultsBlock block = (IntermediateResultsBlock) groupByOperator.nextBlock();
    return block.getAggregationGroupByResult();
}

Example 64 with BrokerRequest

use of com.linkedin.pinot.common.request.BrokerRequest in project pinot by linkedin.

the class TransformExpressionParserTest method testGroupBy.

/**
   * This test ensures that both column names and expressions can be parsed and populated into
   * broker request correctly.
   */
@Test
public void testGroupBy() {
    String[] expectedGroupByColumns = new String[] { "col1", "col2", "col3" };
    String[] expectedGroupByExpressions = new String[] { "udf1(col2)", "ud2(udf3(col1, col2), col3)" };
    String query = "select col1 from myTable group by " + StringUtil.join(",", expectedGroupByColumns) + "," + StringUtil.join(",", expectedGroupByExpressions);
    Pql2Compiler compiler = new Pql2Compiler();
    BrokerRequest brokerRequest = compiler.compileToBrokerRequest(query);
    List<String> actualGroupByColumns = brokerRequest.getGroupBy().getColumns();
    for (int i = 0; i < actualGroupByColumns.size(); i++) {
        Assert.assertEquals(actualGroupByColumns.get(i), expectedGroupByColumns[i]);
    }
    // Group by expression contains columns as well as expressions.
    List<String> actualGroupByExpressions = brokerRequest.getGroupBy().getExpressions();
    for (int i = 0; i < expectedGroupByColumns.length; i++) {
        Assert.assertEquals(actualGroupByExpressions.get(i), expectedGroupByColumns[i]);
    }
    for (int i = 0; i < expectedGroupByExpressions.length; i++) {
        int expressionIndex = i + expectedGroupByColumns.length;
        Assert.assertEquals(actualGroupByExpressions.get(expressionIndex), expectedGroupByExpressions[i].replaceAll("\\s", ""));
    }
}

Example 65 with BrokerRequest

use of com.linkedin.pinot.common.request.BrokerRequest in project pinot by linkedin.

the class QueriesSentinelTest method testTrace.

@Test
public void testTrace() throws Exception {
    String query = "select count(*) from testTable where column1='186154188'";
    LOGGER.info("running  : " + query);
    final Map<ServerInstance, DataTable> instanceResponseMap = new HashMap<ServerInstance, DataTable>();
    final BrokerRequest brokerRequest = REQUEST_COMPILER.compileToBrokerRequest(query);
    //
    brokerRequest.setEnableTrace(true);
    InstanceRequest instanceRequest = new InstanceRequest(1, brokerRequest);
    // TODO: add trace settings consistency
    instanceRequest.setEnableTrace(true);
    instanceRequest.setSearchSegments(new ArrayList<String>());
    instanceRequest.getSearchSegments().add(segmentName);
    QueryRequest queryRequest = new QueryRequest(instanceRequest, TableDataManagerProvider.getServerMetrics());
    final DataTable instanceResponse = QUERY_EXECUTOR.processQuery(queryRequest, queryRunners);
    instanceResponseMap.clear();
    instanceResponseMap.put(new ServerInstance("localhost:0000"), instanceResponse);
    final BrokerResponseNative brokerResponse = REDUCE_SERVICE.reduceOnDataTable(brokerRequest, instanceResponseMap);
    LOGGER.info("BrokerResponse is " + brokerResponse.getAggregationResults().get(0));
    //
    LOGGER.info("TraceInfo is " + brokerResponse.getTraceInfo());
}