Examples with Aggregate org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Aggregate used on opensource projects

Example 26 with Aggregate

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Aggregate in project beam by apache.

the class AggregateScanConverter method convertAggCall.

private AggregateCall convertAggCall(ResolvedComputedColumn computedColumn, int columnRefOff, int groupCount, RelNode input) {
    ResolvedAggregateFunctionCall aggregateFunctionCall = (ResolvedAggregateFunctionCall) computedColumn.getExpr();
    // Reject AVG(INT64)
    if (aggregateFunctionCall.getFunction().getName().equals("avg")) {
        FunctionSignature signature = aggregateFunctionCall.getSignature();
        if (signature.getFunctionArgumentList().get(0).getType().getKind().equals(TypeKind.TYPE_INT64)) {
            throw new UnsupportedOperationException(AVG_ILLEGAL_LONG_INPUT_TYPE);
        }
    }
    // Reject aggregation DISTINCT
    if (aggregateFunctionCall.getDistinct()) {
        throw new UnsupportedOperationException("Does not support " + aggregateFunctionCall.getFunction().getSqlName() + " DISTINCT. 'SELECT DISTINCT' syntax could be used to deduplicate before" + " aggregation.");
    }
    final SqlAggFunction sqlAggFunction;
    if (aggregateFunctionCall.getFunction().getGroup().equals(BeamZetaSqlCatalog.USER_DEFINED_JAVA_AGGREGATE_FUNCTIONS)) {
        // Create a new operator for user-defined functions.
        SqlReturnTypeInference typeInference = x -> ZetaSqlCalciteTranslationUtils.toCalciteType(aggregateFunctionCall.getFunction().getSignatureList().get(0).getResultType().getType(), // TODO(BEAM-9514) set nullable=true
        false, getCluster().getRexBuilder());
        UdafImpl<?, ?, ?> impl = new UdafImpl<>(getExpressionConverter().userFunctionDefinitions.javaAggregateFunctions().get(aggregateFunctionCall.getFunction().getNamePath()));
        sqlAggFunction = SqlOperators.createUdafOperator(aggregateFunctionCall.getFunction().getName(), typeInference, impl);
    } else {
        // Look up builtin functions in SqlOperatorMappingTable.
        sqlAggFunction = (SqlAggFunction) SqlOperatorMappingTable.create(aggregateFunctionCall);
        if (sqlAggFunction == null) {
            throw new UnsupportedOperationException("Does not support ZetaSQL aggregate function: " + aggregateFunctionCall.getFunction().getName());
        }
    }
    List<Integer> argList = new ArrayList<>();
    ResolvedAggregateFunctionCall expr = ((ResolvedAggregateFunctionCall) computedColumn.getExpr());
    List<ZetaSQLResolvedNodeKind.ResolvedNodeKind> resolvedNodeKinds = Arrays.asList(RESOLVED_CAST, RESOLVED_COLUMN_REF, RESOLVED_GET_STRUCT_FIELD);
    for (int i = 0; i < expr.getArgumentList().size(); i++) {
        // Throw an error if aggregate function's input isn't either a ColumnRef or a cast(ColumnRef).
        // TODO: is there a general way to handle aggregation calls conversion?
        ZetaSQLResolvedNodeKind.ResolvedNodeKind resolvedNodeKind = expr.getArgumentList().get(i).nodeKind();
        if (i == 0 && resolvedNodeKinds.contains(resolvedNodeKind)) {
            argList.add(columnRefOff);
        } else if (i > 0 && resolvedNodeKind == RESOLVED_LITERAL) {
            continue;
        } else {
            throw new UnsupportedOperationException("Aggregate function only accepts Column Reference or CAST(Column Reference) as the first argument and " + "Literals as subsequent arguments as its inputs");
        }
    }
    String aggName = getTrait().resolveAlias(computedColumn.getColumn());
    return AggregateCall.create(sqlAggFunction, false, false, false, argList, -1, null, RelCollations.EMPTY, groupCount, input, // When we pass null as the return type, Calcite infers it for us.
    null, aggName);
}

Example 27 with Aggregate

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Aggregate in project beam by apache.

the class BeamBasicAggregationRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    Aggregate aggregate = call.rel(0);
    RelNode relNode = call.rel(1);
    if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
        return;
    }
    if (relNode instanceof Project || relNode instanceof Calc || relNode instanceof Filter) {
        if (isWindowed(relNode) || hasWindowedParents(relNode)) {
            // This case is expected to get handled by the 'BeamAggregationRule'
            return;
        }
    }
    RelNode newTableScan = relNode.copy(relNode.getTraitSet(), relNode.getInputs());
    call.transformTo(new BeamAggregationRel(aggregate.getCluster(), aggregate.getTraitSet().replace(BeamLogicalConvention.INSTANCE), convert(newTableScan, newTableScan.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList(), null, -1));
}

Example 28 with Aggregate

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Aggregate in project beam by apache.

the class SqlCreateFunction method execute.

@Override
public void execute(CalcitePrepare.Context context) {
    final Pair<CalciteSchema, String> pair = SqlDdlNodes.schema(context, true, functionName);
    SchemaPlus schema = pair.left.plus();
    String lastName = pair.right;
    if (!schema.getFunctions(lastName).isEmpty()) {
        throw SqlUtil.newContextException(functionName.getParserPosition(), RESOURCE.internal(String.format("Function %s is already defined.", lastName)));
    }
    JavaUdfLoader udfLoader = new JavaUdfLoader();
    // TODO(BEAM-12355) Support qualified function names.
    List<String> functionPath = ImmutableList.of(lastName);
    if (!(jarPath instanceof SqlCharStringLiteral)) {
        throw SqlUtil.newContextException(jarPath.getParserPosition(), RESOURCE.internal("Jar path is not instanceof SqlCharStringLiteral."));
    }
    String unquotedJarPath = ((SqlCharStringLiteral) jarPath).getNlsString().getValue();
    if (isAggregate) {
        // Try loading the aggregate function just to make sure it exists. LazyAggregateCombineFn will
        // need to fetch it again at runtime.
        udfLoader.loadAggregateFunction(functionPath, unquotedJarPath);
        LazyAggregateCombineFn<?, ?, ?> combineFn = new LazyAggregateCombineFn<>(functionPath, unquotedJarPath);
        schema.add(lastName, combineFn.getUdafImpl());
    } else {
        ScalarFn scalarFn = udfLoader.loadScalarFunction(functionPath, unquotedJarPath);
        Method method = ScalarFnReflector.getApplyMethod(scalarFn);
        Function function = ScalarFunctionImpl.create(method, unquotedJarPath);
        schema.add(lastName, function);
    }
}

Example 29 with Aggregate

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Aggregate in project druid by druid-io.

the class DruidQuery method computeGrouping.

@Nonnull
private static Grouping computeGrouping(final PartialDruidQuery partialQuery, final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final boolean finalizeAggregations) {
    final Aggregate aggregate = Preconditions.checkNotNull(partialQuery.getAggregate(), "aggregate");
    final Project aggregateProject = partialQuery.getAggregateProject();
    final List<DimensionExpression> dimensions = computeDimensions(partialQuery, plannerContext, rowSignature, virtualColumnRegistry);
    final Subtotals subtotals = computeSubtotals(partialQuery, rowSignature);
    final List<Aggregation> aggregations = computeAggregations(partialQuery, plannerContext, rowSignature, virtualColumnRegistry, rexBuilder, finalizeAggregations);
    final RowSignature aggregateRowSignature = RowSignatures.fromRelDataType(ImmutableList.copyOf(Iterators.concat(dimensions.stream().map(DimensionExpression::getOutputName).iterator(), aggregations.stream().map(Aggregation::getOutputName).iterator())), aggregate.getRowType());
    final DimFilter havingFilter = computeHavingFilter(partialQuery, plannerContext, aggregateRowSignature);
    final Grouping grouping = Grouping.create(dimensions, subtotals, aggregations, havingFilter, aggregateRowSignature);
    if (aggregateProject == null) {
        return grouping;
    } else {
        return grouping.applyProject(plannerContext, aggregateProject);
    }
}

Example 30 with Aggregate

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Aggregate in project druid by druid-io.

the class DruidQuery method computeDimensions.

/**
 * Returns dimensions corresponding to {@code aggregate.getGroupSet()}, in the same order.
 *
 * @param partialQuery          partial query
 * @param plannerContext        planner context
 * @param rowSignature          source row signature
 * @param virtualColumnRegistry re-usable virtual column references
 *
 * @return dimensions
 *
 * @throws CannotBuildQueryException if dimensions cannot be computed
 */
private static List<DimensionExpression> computeDimensions(final PartialDruidQuery partialQuery, final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry) {
    final Aggregate aggregate = Preconditions.checkNotNull(partialQuery.getAggregate());
    final List<DimensionExpression> dimensions = new ArrayList<>();
    final String outputNamePrefix = Calcites.findUnusedPrefixForDigits("d", rowSignature.getColumnNames());
    int outputNameCounter = 0;
    for (int i : aggregate.getGroupSet()) {
        // Dimension might need to create virtual columns. Avoid giving it a name that would lead to colliding columns.
        final RexNode rexNode = Expressions.fromFieldAccess(rowSignature, partialQuery.getSelectProject(), i);
        final DruidExpression druidExpression = Expressions.toDruidExpression(plannerContext, rowSignature, rexNode);
        if (druidExpression == null) {
            throw new CannotBuildQueryException(aggregate, rexNode);
        }
        final RelDataType dataType = rexNode.getType();
        final ColumnType outputType = Calcites.getColumnTypeForRelDataType(dataType);
        if (Types.isNullOr(outputType, ValueType.COMPLEX)) {
            // Can't group on unknown or COMPLEX types.
            plannerContext.setPlanningError("SQL requires a group-by on a column of type %s that is unsupported.", outputType);
            throw new CannotBuildQueryException(aggregate, rexNode);
        }
        final String dimOutputName = outputNamePrefix + outputNameCounter++;
        if (!druidExpression.isSimpleExtraction()) {
            final String virtualColumn = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(druidExpression, dataType);
            dimensions.add(DimensionExpression.ofVirtualColumn(virtualColumn, dimOutputName, druidExpression, outputType));
        } else {
            dimensions.add(DimensionExpression.ofSimpleColumn(dimOutputName, druidExpression, outputType));
        }
    }
    return dimensions;
}