Search in sources :

Example 6 with ExprMacroTable

use of org.apache.druid.math.expr.ExprMacroTable in project druid by druid-io.

the class ArrayConcatSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(PlannerContext plannerContext, RowSignature rowSignature, VirtualColumnRegistry virtualColumnRegistry, RexBuilder rexBuilder, String name, AggregateCall aggregateCall, Project project, List<Aggregation> existingAggregations, boolean finalizeAggregations) {
    final List<RexNode> arguments = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).collect(Collectors.toList());
    Integer maxSizeBytes = null;
    if (arguments.size() > 1) {
        RexNode maxBytes = arguments.get(1);
        if (!maxBytes.isA(SqlKind.LITERAL)) {
            // maxBytes must be a literal
            return null;
        }
        maxSizeBytes = ((Number) RexLiteral.value(maxBytes)).intValue();
    }
    final DruidExpression arg = Expressions.toDruidExpression(plannerContext, rowSignature, arguments.get(0));
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String fieldName;
    final ColumnType druidType = Calcites.getValueTypeForRelDataTypeFull(aggregateCall.getType());
    if (druidType == null || !druidType.isArray()) {
        // must be an array
        return null;
    }
    final String initialvalue = ExpressionType.fromColumnTypeStrict(druidType).asTypeString() + "[]";
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        VirtualColumn vc = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(plannerContext, arg, druidType);
        fieldName = vc.getOutputName();
    }
    if (aggregateCall.isDistinct()) {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", name), null, null, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable));
    } else {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_concat(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_concat(\"__acc\", \"%s\")", name), null, null, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable));
    }
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) ReturnTypes(org.apache.calcite.sql.type.ReturnTypes) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) Optionality(org.apache.calcite.util.Optionality) RexNode(org.apache.calcite.rex.RexNode) ExpressionType(org.apache.druid.math.expr.ExpressionType) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) Nullable(javax.annotation.Nullable) ImmutableSet(com.google.common.collect.ImmutableSet) SqlKind(org.apache.calcite.sql.SqlKind) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) InferTypes(org.apache.calcite.sql.type.InferTypes) RexBuilder(org.apache.calcite.rex.RexBuilder) RexLiteral(org.apache.calcite.rex.RexLiteral) VirtualColumn(org.apache.druid.segment.VirtualColumn) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) StringUtils(org.apache.druid.java.util.common.StringUtils) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Collectors(java.util.stream.Collectors) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) List(java.util.List) RowSignature(org.apache.druid.segment.column.RowSignature) OperandTypes(org.apache.calcite.sql.type.OperandTypes) ColumnType(org.apache.druid.segment.column.ColumnType) AggregateCall(org.apache.calcite.rel.core.AggregateCall) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) Calcites(org.apache.druid.sql.calcite.planner.Calcites) Expressions(org.apache.druid.sql.calcite.expression.Expressions) ColumnType(org.apache.druid.segment.column.ColumnType) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) VirtualColumn(org.apache.druid.segment.VirtualColumn) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Example 7 with ExprMacroTable

use of org.apache.druid.math.expr.ExprMacroTable in project druid by druid-io.

the class AvgSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final String name, final AggregateCall aggregateCall, final Project project, final List<Aggregation> existingAggregations, final boolean finalizeAggregations) {
    final List<DruidExpression> arguments = Aggregations.getArgumentsForSimpleAggregator(plannerContext, rowSignature, aggregateCall, project);
    if (arguments == null) {
        return null;
    }
    final String countName = Calcites.makePrefixedName(name, "count");
    final AggregatorFactory count = CountSqlAggregator.createCountAggregatorFactory(countName, plannerContext, rowSignature, virtualColumnRegistry, rexBuilder, aggregateCall, project);
    final String fieldName;
    final String expression;
    final DruidExpression arg = Iterables.getOnlyElement(arguments);
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final ValueType sumType;
    // Use 64-bit sum regardless of the type of the AVG aggregator.
    if (SqlTypeName.INT_TYPES.contains(aggregateCall.getType().getSqlTypeName())) {
        sumType = ValueType.LONG;
    } else {
        sumType = ValueType.DOUBLE;
    }
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
        expression = null;
    } else {
        // if the filter or anywhere else defined a virtual column for us, re-use it
        final RexNode resolutionArg = Expressions.fromFieldAccess(rowSignature, project, Iterables.getOnlyElement(aggregateCall.getArgList()));
        String vc = virtualColumnRegistry.getVirtualColumnByExpression(arg, resolutionArg.getType());
        fieldName = vc != null ? vc : null;
        expression = vc != null ? null : arg.getExpression();
    }
    final String sumName = Calcites.makePrefixedName(name, "sum");
    final AggregatorFactory sum = SumSqlAggregator.createSumAggregatorFactory(sumType, sumName, fieldName, expression, macroTable);
    return Aggregation.create(ImmutableList.of(sum, count), new ArithmeticPostAggregator(name, "quotient", ImmutableList.of(new FieldAccessPostAggregator(null, sumName), new FieldAccessPostAggregator(null, countName))));
}
Also used : ArithmeticPostAggregator(org.apache.druid.query.aggregation.post.ArithmeticPostAggregator) FieldAccessPostAggregator(org.apache.druid.query.aggregation.post.FieldAccessPostAggregator) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ValueType(org.apache.druid.segment.column.ValueType) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Example 8 with ExprMacroTable

use of org.apache.druid.math.expr.ExprMacroTable in project druid by druid-io.

the class MultiColumnSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final String name, final AggregateCall aggregateCall, final Project project, final List<Aggregation> existingAggregations, final boolean finalizeAggregations) {
    if (aggregateCall.isDistinct()) {
        return null;
    }
    final List<DruidExpression> arguments = Aggregations.getArgumentsForSimpleAggregator(plannerContext, rowSignature, aggregateCall, project);
    if (arguments == null) {
        return null;
    }
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final List<FieldInfo> fieldInfoList = new ArrayList<>();
    // Convert arguments to concise field information
    for (DruidExpression argument : arguments) {
        if (argument.isDirectColumnAccess()) {
            fieldInfoList.add(FieldInfo.fromFieldName(argument.getDirectColumn()));
        } else {
            fieldInfoList.add(FieldInfo.fromExpression(argument.getExpression()));
        }
    }
    Preconditions.checkArgument(!fieldInfoList.isEmpty(), "FieldInfoList should not be empty");
    return getAggregation(name, aggregateCall, macroTable, fieldInfoList);
}
Also used : DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ArrayList(java.util.ArrayList) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) Nullable(javax.annotation.Nullable)

Example 9 with ExprMacroTable

use of org.apache.druid.math.expr.ExprMacroTable in project druid by druid-io.

the class JoinAndLookupBenchmark method setup.

@Setup()
public void setup() throws IOException {
    tmpDir = FileUtils.createTempDir();
    ColumnConfig columnConfig = () -> columnCacheSizeBytes;
    index = JoinTestHelper.createFactIndexBuilder(columnConfig, tmpDir, rows).buildMMappedIndex();
    final String prefix = "c.";
    baseSegment = new QueryableIndexSegment(index, SegmentId.dummy("join"));
    List<JoinableClause> joinableClausesLookupStringKey = ImmutableList.of(new JoinableClause(prefix, LookupJoinable.wrap(JoinTestHelper.createCountryIsoCodeToNameLookup()), JoinType.LEFT, JoinConditionAnalysis.forExpression(StringUtils.format("countryIsoCode == \"%sk\"", prefix), prefix, ExprMacroTable.nil())));
    JoinFilterPreAnalysis preAnalysisLookupStringKey = JoinFilterAnalyzer.computeJoinFilterPreAnalysis(new JoinFilterPreAnalysisKey(new JoinFilterRewriteConfig(false, false, false, false, 0), joinableClausesLookupStringKey, VirtualColumns.EMPTY, null));
    hashJoinLookupStringKeySegment = new HashJoinSegment(ReferenceCountingSegment.wrapRootGenerationSegment(baseSegment), null, joinableClausesLookupStringKey, preAnalysisLookupStringKey);
    List<JoinableClause> joinableClausesLookupLongKey = ImmutableList.of(new JoinableClause(prefix, LookupJoinable.wrap(JoinTestHelper.createCountryIsoCodeToNameLookup()), JoinType.LEFT, JoinConditionAnalysis.forExpression(StringUtils.format("countryIsoCode == \"%sk\"", prefix), prefix, ExprMacroTable.nil())));
    JoinFilterPreAnalysis preAnalysisLookupLongKey = JoinFilterAnalyzer.computeJoinFilterPreAnalysis(new JoinFilterPreAnalysisKey(new JoinFilterRewriteConfig(false, false, false, false, 0), joinableClausesLookupLongKey, VirtualColumns.EMPTY, null));
    hashJoinLookupLongKeySegment = new HashJoinSegment(ReferenceCountingSegment.wrapRootGenerationSegment(baseSegment), null, joinableClausesLookupLongKey, preAnalysisLookupLongKey);
    List<JoinableClause> joinableClausesIndexedTableStringKey = ImmutableList.of(new JoinableClause(prefix, new IndexedTableJoinable(JoinTestHelper.createCountriesIndexedTable()), JoinType.LEFT, JoinConditionAnalysis.forExpression(StringUtils.format("countryIsoCode == \"%scountryIsoCode\"", prefix), prefix, ExprMacroTable.nil())));
    JoinFilterPreAnalysis preAnalysisIndexedStringKey = JoinFilterAnalyzer.computeJoinFilterPreAnalysis(new JoinFilterPreAnalysisKey(new JoinFilterRewriteConfig(false, false, false, false, 0), joinableClausesLookupLongKey, VirtualColumns.EMPTY, null));
    hashJoinIndexedTableStringKeySegment = new HashJoinSegment(ReferenceCountingSegment.wrapRootGenerationSegment(baseSegment), null, joinableClausesIndexedTableStringKey, preAnalysisIndexedStringKey);
    List<JoinableClause> joinableClausesIndexedTableLongKey = ImmutableList.of(new JoinableClause(prefix, new IndexedTableJoinable(JoinTestHelper.createCountriesIndexedTable()), JoinType.LEFT, JoinConditionAnalysis.forExpression(StringUtils.format("countryNumber == \"%scountryNumber\"", prefix), prefix, ExprMacroTable.nil())));
    JoinFilterPreAnalysis preAnalysisIndexedLongKey = JoinFilterAnalyzer.computeJoinFilterPreAnalysis(new JoinFilterPreAnalysisKey(new JoinFilterRewriteConfig(false, false, false, false, 0), joinableClausesIndexedTableLongKey, VirtualColumns.EMPTY, null));
    hashJoinIndexedTableLongKeySegment = new HashJoinSegment(ReferenceCountingSegment.wrapRootGenerationSegment(baseSegment), null, joinableClausesIndexedTableLongKey, preAnalysisIndexedLongKey);
    final Map<String, String> countryCodeToNameMap = JoinTestHelper.createCountryIsoCodeToNameLookup().getMap();
    final Map<String, String> countryNumberToNameMap = JoinTestHelper.createCountryNumberToNameLookup().getMap();
    final ExprMacroTable exprMacroTable = new ExprMacroTable(ImmutableList.of(new LookupExprMacro(new LookupExtractorFactoryContainerProvider() {

        @Override
        public Set<String> getAllLookupNames() {
            return ImmutableSet.of(LOOKUP_COUNTRY_CODE_TO_NAME, LOOKUP_COUNTRY_NUMBER_TO_NAME);
        }

        @Override
        public Optional<LookupExtractorFactoryContainer> get(String lookupName) {
            if (LOOKUP_COUNTRY_CODE_TO_NAME.equals(lookupName)) {
                return Optional.of(new LookupExtractorFactoryContainer("0", new MapLookupExtractorFactory(countryCodeToNameMap, false)));
            } else if (LOOKUP_COUNTRY_NUMBER_TO_NAME.equals(lookupName)) {
                return Optional.of(new LookupExtractorFactoryContainer("0", new MapLookupExtractorFactory(countryNumberToNameMap, false)));
            } else {
                return Optional.empty();
            }
        }
    })));
    lookupVirtualColumns = VirtualColumns.create(ImmutableList.of(new ExpressionVirtualColumn(LOOKUP_COUNTRY_CODE_TO_NAME, "lookup(countryIsoCode, '" + LOOKUP_COUNTRY_CODE_TO_NAME + "')", ColumnType.STRING, exprMacroTable), new ExpressionVirtualColumn(LOOKUP_COUNTRY_NUMBER_TO_NAME, "lookup(countryNumber, '" + LOOKUP_COUNTRY_NUMBER_TO_NAME + "')", ColumnType.STRING, exprMacroTable)));
}
Also used : QueryableIndexSegment(org.apache.druid.segment.QueryableIndexSegment) ImmutableSet(com.google.common.collect.ImmutableSet) Set(java.util.Set) ColumnConfig(org.apache.druid.segment.column.ColumnConfig) Optional(java.util.Optional) MapLookupExtractorFactory(org.apache.druid.query.lookup.MapLookupExtractorFactory) JoinFilterPreAnalysisKey(org.apache.druid.segment.join.filter.JoinFilterPreAnalysisKey) LookupExprMacro(org.apache.druid.query.expression.LookupExprMacro) JoinFilterRewriteConfig(org.apache.druid.segment.join.filter.rewrite.JoinFilterRewriteConfig) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) JoinableClause(org.apache.druid.segment.join.JoinableClause) HashJoinSegment(org.apache.druid.segment.join.HashJoinSegment) ExpressionVirtualColumn(org.apache.druid.segment.virtual.ExpressionVirtualColumn) JoinFilterPreAnalysis(org.apache.druid.segment.join.filter.JoinFilterPreAnalysis) IndexedTableJoinable(org.apache.druid.segment.join.table.IndexedTableJoinable) LookupExtractorFactoryContainerProvider(org.apache.druid.query.lookup.LookupExtractorFactoryContainerProvider) LookupExtractorFactoryContainer(org.apache.druid.query.lookup.LookupExtractorFactoryContainer) Setup(org.openjdk.jmh.annotations.Setup)

Example 10 with ExprMacroTable

use of org.apache.druid.math.expr.ExprMacroTable in project druid by druid-io.

the class SleepSqlTest method createMacroTable.

@Override
public ExprMacroTable createMacroTable() {
    final List<ExprMacro> exprMacros = new ArrayList<>();
    for (Class<? extends ExprMacroTable.ExprMacro> clazz : ExpressionModule.EXPR_MACROS) {
        exprMacros.add(CalciteTests.INJECTOR.getInstance(clazz));
    }
    exprMacros.add(CalciteTests.INJECTOR.getInstance(LookupExprMacro.class));
    exprMacros.add(new SleepExprMacro());
    return new ExprMacroTable(exprMacros);
}
Also used : ArrayList(java.util.ArrayList) LookupExprMacro(org.apache.druid.query.expression.LookupExprMacro) SleepExprMacro(org.apache.druid.query.expressions.SleepExprMacro) ExprMacro(org.apache.druid.math.expr.ExprMacroTable.ExprMacro) LookupExprMacro(org.apache.druid.query.expression.LookupExprMacro) SleepExprMacro(org.apache.druid.query.expressions.SleepExprMacro) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable)

Aggregations

ExprMacroTable (org.apache.druid.math.expr.ExprMacroTable)16 Nullable (javax.annotation.Nullable)7 DruidExpression (org.apache.druid.sql.calcite.expression.DruidExpression)7 ImmutableSet (com.google.common.collect.ImmutableSet)5 DruidOperatorTable (org.apache.druid.sql.calcite.planner.DruidOperatorTable)5 PlannerConfig (org.apache.druid.sql.calcite.planner.PlannerConfig)5 DruidSchemaCatalog (org.apache.druid.sql.calcite.schema.DruidSchemaCatalog)5 ArrayList (java.util.ArrayList)4 List (java.util.List)4 Collectors (java.util.stream.Collectors)4 AggregateCall (org.apache.calcite.rel.core.AggregateCall)4 Project (org.apache.calcite.rel.core.Project)4 RexBuilder (org.apache.calcite.rex.RexBuilder)4 RexNode (org.apache.calcite.rex.RexNode)4 SqlAggFunction (org.apache.calcite.sql.SqlAggFunction)4 SqlFunctionCategory (org.apache.calcite.sql.SqlFunctionCategory)4 SqlKind (org.apache.calcite.sql.SqlKind)4 InferTypes (org.apache.calcite.sql.type.InferTypes)4 OperandTypes (org.apache.calcite.sql.type.OperandTypes)4 Optionality (org.apache.calcite.util.Optionality)4