Examples with PinotQueryGenerator com.facebook.presto.pinot.query.PinotQueryGenerator used on opensource projects

Example 1 with PinotQueryGenerator

use of com.facebook.presto.pinot.query.PinotQueryGenerator in project presto by prestodb.

the class TestPinotSplitManager method testSegmentSplitsHelperNoFilter.

private void testSegmentSplitsHelperNoFilter(PinotTableHandle table, int segmentsPerSplit, int expectedNumSplits, boolean expectFilter) {
    PinotConfig pinotConfig = new PinotConfig().setForbidBrokerQueries(false);
    SessionHolder sessionHolder = new SessionHolder(pinotConfig);
    PlanBuilder planBuilder = createPlanBuilder(sessionHolder);
    PlanNode plan = tableScan(planBuilder, table, regionId, city, fare, secondsSinceEpoch);
    PinotQueryGenerator.PinotQueryGeneratorResult pinotQueryGeneratorResult = new PinotQueryGenerator(pinotConfig, functionAndTypeManager, functionAndTypeManager, standardFunctionResolution).generate(plan, sessionHolder.getConnectorSession()).get();
    List<PinotColumnHandle> expectedHandles = ImmutableList.copyOf(pinotQueryGeneratorResult.getContext().getAssignments(pinotQueryGeneratorResult.getGeneratedPinotQuery().getFormat() == PinotQueryGenerator.PinotQueryFormat.SQL).values());
    PinotQueryGenerator.GeneratedPinotQuery generatedPql = pinotQueryGeneratorResult.getGeneratedPinotQuery();
    PinotTableHandle pinotTableHandle = new PinotTableHandle(table.getConnectorId(), table.getSchemaName(), table.getTableName(), Optional.of(false), Optional.of(expectedHandles), Optional.of(generatedPql));
    List<PinotSplit> splits = getSplitsHelper(pinotTableHandle, segmentsPerSplit, false);
    assertSplits(splits, expectedNumSplits, SEGMENT);
    splits.forEach(s -> assertSegmentSplitWellFormed(s, expectFilter));
}

Example 2 with PinotQueryGenerator

use of com.facebook.presto.pinot.query.PinotQueryGenerator in project presto by prestodb.

the class TestPinotSplitManager method testBrokerTopNLarge.

private void testBrokerTopNLarge(PinotTableHandle table, int sessionTopNLarge, int configTopNLarge, boolean useSql) {
    String topNLimitKeyword = useSql ? "LIMIT " : "TOP ";
    PinotConfig pinotConfig = new PinotConfig().setUsePinotSqlForBrokerQueries(useSql).setTopNLarge(configTopNLarge);
    SessionHolder sessionHolder = new SessionHolder(pinotConfig);
    ConnectorSession session = createSessionWithTopNLarge(sessionTopNLarge, pinotConfig);
    PlanBuilder planBuilder = createPlanBuilder(sessionHolder);
    PlanNode tableScanNode = tableScan(planBuilder, table, regionId, city, fare, secondsSinceEpoch);
    AggregationNode aggregationNode = planBuilder.aggregation(aggregationNodeBuilder -> aggregationNodeBuilder.source(tableScanNode).singleGroupingSet(variable("city"), variable("regionid")).addAggregation(planBuilder.variable("sum_fare"), getRowExpression("sum(fare)", sessionHolder)).addAggregation(planBuilder.variable("count_regionid"), getRowExpression("count(regionid)", sessionHolder)));
    PinotQueryGenerator.PinotQueryGeneratorResult pinotQueryGeneratorResult = new PinotQueryGenerator(pinotConfig, functionAndTypeManager, functionAndTypeManager, standardFunctionResolution).generate(aggregationNode, session).get();
    String[] limits = pinotQueryGeneratorResult.getGeneratedPinotQuery().getQuery().split(topNLimitKeyword);
    assertEquals(Integer.parseInt(limits[1]), sessionTopNLarge);
    aggregationNode = planBuilder.aggregation(aggregationNodeBuilder -> aggregationNodeBuilder.source(tableScanNode).singleGroupingSet(variable("city"), variable("regionid")).addAggregation(planBuilder.variable("sum_fare"), getRowExpression("sum(fare)", sessionHolder)).addAggregation(planBuilder.variable("count_regionid"), getRowExpression("count(regionid)", sessionHolder)));
    pinotQueryGeneratorResult = new PinotQueryGenerator(pinotConfig, functionAndTypeManager, functionAndTypeManager, standardFunctionResolution).generate(aggregationNode, sessionHolder.getConnectorSession()).get();
    limits = pinotQueryGeneratorResult.getGeneratedPinotQuery().getQuery().split(topNLimitKeyword);
    assertEquals(Integer.parseInt(limits[1]), configTopNLarge);
}

Example 3 with PinotQueryGenerator

use of com.facebook.presto.pinot.query.PinotQueryGenerator in project presto by prestodb.

the class TestPinotSplitManager method testSegmentSplitsHelperWithFilter.

private void testSegmentSplitsHelperWithFilter(PinotTableHandle table, int segmentsPerSplit, int expectedNumSplits) {
    PinotConfig pinotConfig = new PinotConfig().setForbidBrokerQueries(false);
    SessionHolder sessionHolder = new SessionHolder(pinotConfig);
    PlanBuilder planBuilder = createPlanBuilder(sessionHolder);
    PlanNode plan = filter(planBuilder, tableScan(planBuilder, table, regionId, city, fare, secondsSinceEpoch), getRowExpression("city = 'Boston'", sessionHolder));
    PinotQueryGenerator.PinotQueryGeneratorResult pinotQueryGeneratorResult = new PinotQueryGenerator(pinotConfig, functionAndTypeManager, functionAndTypeManager, standardFunctionResolution).generate(plan, sessionHolder.getConnectorSession()).get();
    List<PinotColumnHandle> expectedHandles = ImmutableList.copyOf(pinotQueryGeneratorResult.getContext().getAssignments(pinotQueryGeneratorResult.getGeneratedPinotQuery().getFormat() == PinotQueryGenerator.PinotQueryFormat.SQL).values());
    PinotQueryGenerator.GeneratedPinotQuery generatedPql = pinotQueryGeneratorResult.getGeneratedPinotQuery();
    PinotTableHandle pinotTableHandle = new PinotTableHandle(table.getConnectorId(), table.getSchemaName(), table.getTableName(), Optional.of(false), Optional.of(expectedHandles), Optional.of(generatedPql));
    List<PinotSplit> splits = getSplitsHelper(pinotTableHandle, segmentsPerSplit, false);
    assertSplits(splits, expectedNumSplits, SEGMENT);
    splits.forEach(s -> assertSegmentSplitWellFormed(s, true));
}

Example 4 with PinotQueryGenerator

use of com.facebook.presto.pinot.query.PinotQueryGenerator in project presto by prestodb.

the class TestPinotSplitManager method testSegmentLimitLarge.

private void testSegmentLimitLarge(PinotTableHandle table, int sessionLimitLarge, int configLimitLarge, boolean useSql) {
    PinotConfig pinotConfig = new PinotConfig().setUsePinotSqlForBrokerQueries(useSql).setLimitLargeForSegment(configLimitLarge);
    SessionHolder sessionHolder = new SessionHolder(pinotConfig);
    ConnectorSession session = createSessionWithLimitLarge(sessionLimitLarge, pinotConfig);
    PlanBuilder planBuilder = createPlanBuilder(sessionHolder);
    PlanNode plan = tableScan(planBuilder, table, regionId, city, fare, secondsSinceEpoch);
    PinotQueryGenerator.PinotQueryGeneratorResult pinotQueryGeneratorResult = new PinotQueryGenerator(pinotConfig, functionAndTypeManager, functionAndTypeManager, standardFunctionResolution).generate(plan, session).get();
    String[] limits = pinotQueryGeneratorResult.getGeneratedPinotQuery().getQuery().split("LIMIT ");
    assertEquals(Integer.parseInt(limits[1]), sessionLimitLarge);
    plan = tableScan(planBuilder, table, regionId, city, fare, secondsSinceEpoch);
    pinotQueryGeneratorResult = new PinotQueryGenerator(pinotConfig, functionAndTypeManager, functionAndTypeManager, standardFunctionResolution).generate(plan, sessionHolder.getConnectorSession()).get();
    limits = pinotQueryGeneratorResult.getGeneratedPinotQuery().getQuery().split("LIMIT ");
    assertEquals(Integer.parseInt(limits[1]), configLimitLarge);
}