Examples with Pair org.apache.calcite.util.Pair used on opensource projects

Example 76 with Pair

use of org.apache.calcite.util.Pair in project beam by apache.

the class BeamJoinRel method extractJoinRexNodes.

static List<Pair<RexNode, RexNode>> extractJoinRexNodes(RexNode condition) {
    // or it's a JOIN ON false because: condition == false
    if (condition instanceof RexLiteral) {
        throw new UnsupportedOperationException("CROSS JOIN, JOIN ON FALSE is not supported!");
    }
    RexCall call = (RexCall) condition;
    List<Pair<RexNode, RexNode>> pairs = new ArrayList<>();
    if ("AND".equals(call.getOperator().getName())) {
        List<RexNode> operands = call.getOperands();
        for (RexNode rexNode : operands) {
            Pair<RexNode, RexNode> pair = extractJoinPairOfRexNodes((RexCall) rexNode);
            pairs.add(pair);
        }
    } else if ("=".equals(call.getOperator().getName())) {
        pairs.add(extractJoinPairOfRexNodes(call));
    } else {
        throw new UnsupportedOperationException("Operator " + call.getOperator().getName() + " is not supported in join condition");
    }
    return pairs;
}

Example 77 with Pair

use of org.apache.calcite.util.Pair in project hazelcast by hazelcast.

the class HazelcastSqlToRelConverter method constructComparisons.

/**
 * Constructs comparisons between
 * left-hand operand (as a rule, SqlIdentifier) and right-hand list.
 */
private List<RexNode> constructComparisons(Blackboard bb, List<RexNode> leftKeys, SqlNodeList valuesList) {
    final List<RexNode> comparisons = new ArrayList<>();
    for (SqlNode rightValues : valuesList) {
        RexNode rexComparison;
        final SqlOperator comparisonOp = SqlStdOperatorTable.EQUALS;
        if (leftKeys.size() == 1) {
            rexComparison = rexBuilder.makeCall(comparisonOp, leftKeys.get(0), ensureSqlType(leftKeys.get(0).getType(), bb.convertExpression(rightValues)));
        } else {
            assert rightValues instanceof SqlCall;
            final SqlBasicCall basicCall = (SqlBasicCall) rightValues;
            assert basicCall.getOperator() instanceof SqlRowOperator && basicCall.operandCount() == leftKeys.size();
            rexComparison = RexUtil.composeConjunction(rexBuilder, Pair.zip(leftKeys, basicCall.getOperandList()).stream().map(pair -> rexBuilder.makeCall(comparisonOp, pair.left, ensureSqlType(pair.left.getType(), bb.convertExpression(pair.right)))).collect(Collectors.toList()));
        }
        comparisons.add(rexComparison);
    }
    return comparisons;
}

Example 78 with Pair

use of org.apache.calcite.util.Pair in project samza by apache.

the class TestProjectTranslator method testTranslate.

@Test
public void testTranslate() throws IOException, ClassNotFoundException {
    // setup mock values to the constructor of FilterTranslator
    LogicalProject mockProject = PowerMockito.mock(LogicalProject.class);
    Context mockContext = mock(Context.class);
    ContainerContext mockContainerContext = mock(ContainerContext.class);
    TranslatorContext mockTranslatorContext = mock(TranslatorContext.class);
    TestMetricsRegistryImpl testMetricsRegistryImpl = new TestMetricsRegistryImpl();
    RelNode mockInput = mock(RelNode.class);
    List<RelNode> inputs = new ArrayList<>();
    inputs.add(mockInput);
    when(mockInput.getId()).thenReturn(1);
    when(mockProject.getId()).thenReturn(2);
    when(mockProject.getInputs()).thenReturn(inputs);
    when(mockProject.getInput()).thenReturn(mockInput);
    RelDataType mockRowType = mock(RelDataType.class);
    when(mockRowType.getFieldCount()).thenReturn(1);
    when(mockProject.getRowType()).thenReturn(mockRowType);
    RexNode mockRexField = mock(RexNode.class);
    List<Pair<RexNode, String>> namedProjects = new ArrayList<>();
    namedProjects.add(Pair.of(mockRexField, "test_field"));
    when(mockProject.getNamedProjects()).thenReturn(namedProjects);
    StreamApplicationDescriptorImpl mockAppDesc = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec<Object, SamzaSqlRelMessage> mockInputOp = mock(OperatorSpec.class);
    MessageStream<SamzaSqlRelMessage> mockStream = new MessageStreamImpl<>(mockAppDesc, mockInputOp);
    when(mockTranslatorContext.getMessageStream(eq(1))).thenReturn(mockStream);
    doAnswer(this.getRegisterMessageStreamAnswer()).when(mockTranslatorContext).registerMessageStream(eq(2), any(MessageStream.class));
    RexToJavaCompiler mockCompiler = mock(RexToJavaCompiler.class);
    when(mockTranslatorContext.getExpressionCompiler()).thenReturn(mockCompiler);
    Expression mockExpr = mock(Expression.class);
    when(mockCompiler.compile(any(), any())).thenReturn(mockExpr);
    when(mockContext.getContainerContext()).thenReturn(mockContainerContext);
    when(mockContainerContext.getContainerMetricsRegistry()).thenReturn(testMetricsRegistryImpl);
    // Apply translate() method to verify that we are getting the correct map operator constructed
    ProjectTranslator projectTranslator = new ProjectTranslator(1);
    projectTranslator.translate(mockProject, LOGICAL_OP_ID, mockTranslatorContext);
    // make sure that context has been registered with LogicFilter and output message streams
    verify(mockTranslatorContext, times(1)).registerRelNode(2, mockProject);
    verify(mockTranslatorContext, times(1)).registerMessageStream(2, this.getRegisteredMessageStream(2));
    when(mockTranslatorContext.getRelNode(2)).thenReturn(mockProject);
    when(mockTranslatorContext.getMessageStream(2)).thenReturn(this.getRegisteredMessageStream(2));
    StreamOperatorSpec projectSpec = (StreamOperatorSpec) Whitebox.getInternalState(this.getRegisteredMessageStream(2), "operatorSpec");
    assertNotNull(projectSpec);
    assertEquals(projectSpec.getOpCode(), OperatorSpec.OpCode.MAP);
    // Verify that the bootstrap() method will establish the context for the map function
    Map<Integer, TranslatorContext> mockContexts = new HashMap<>();
    mockContexts.put(1, mockTranslatorContext);
    when(mockContext.getApplicationTaskContext()).thenReturn(new SamzaSqlApplicationContext(mockContexts));
    projectSpec.getTransformFn().init(mockContext);
    MapFunction mapFn = (MapFunction) Whitebox.getInternalState(projectSpec, "mapFn");
    assertNotNull(mapFn);
    assertEquals(mockTranslatorContext, Whitebox.getInternalState(mapFn, "translatorContext"));
    assertEquals(mockProject, Whitebox.getInternalState(mapFn, "project"));
    assertEquals(mockExpr, Whitebox.getInternalState(mapFn, "expr"));
    // Verify TestMetricsRegistryImpl works with Project
    assertEquals(1, testMetricsRegistryImpl.getGauges().size());
    assertEquals(2, testMetricsRegistryImpl.getGauges().get(LOGICAL_OP_ID).size());
    assertEquals(1, testMetricsRegistryImpl.getCounters().size());
    assertEquals(2, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).size());
    assertEquals(0, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(0, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
    // Calling mapFn.apply() to verify the filter function is correctly applied to the input message
    SamzaSqlRelMessage mockInputMsg = new SamzaSqlRelMessage(new ArrayList<>(), new ArrayList<>(), new SamzaSqlRelMsgMetadata(0L, 0L));
    SamzaSqlExecutionContext executionContext = mock(SamzaSqlExecutionContext.class);
    DataContext dataContext = mock(DataContext.class);
    when(mockTranslatorContext.getExecutionContext()).thenReturn(executionContext);
    when(mockTranslatorContext.getDataContext()).thenReturn(dataContext);
    Object[] result = new Object[1];
    final Object mockFieldObj = new Object();
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = mockFieldObj;
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    SamzaSqlRelMessage retMsg = (SamzaSqlRelMessage) mapFn.apply(mockInputMsg);
    assertEquals(retMsg.getSamzaSqlRelRecord().getFieldNames(), Collections.singletonList("test_field"));
    assertEquals(retMsg.getSamzaSqlRelRecord().getFieldValues(), Collections.singletonList(mockFieldObj));
    // Verify mapFn.apply() updates the TestMetricsRegistryImpl metrics
    assertEquals(1, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(1, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
}

Example 79 with Pair

use of org.apache.calcite.util.Pair in project flink by apache.

the class HiveParserUtils method projectNonColumnEquiConditions.

/**
 * Push any equi join conditions that are not column references as Projections on top of the
 * children.
 */
public static RexNode projectNonColumnEquiConditions(RelFactories.ProjectFactory factory, RelNode[] inputRels, List<RexNode> leftJoinKeys, List<RexNode> rightJoinKeys, int systemColCount, List<Integer> leftKeys, List<Integer> rightKeys) {
    RelNode leftRel = inputRels[0];
    RelNode rightRel = inputRels[1];
    RexBuilder rexBuilder = leftRel.getCluster().getRexBuilder();
    RexNode outJoinCond = null;
    int origLeftInputSize = leftRel.getRowType().getFieldCount();
    int origRightInputSize = rightRel.getRowType().getFieldCount();
    List<RexNode> newLeftFields = new ArrayList<>();
    List<String> newLeftFieldNames = new ArrayList<>();
    List<RexNode> newRightFields = new ArrayList<>();
    List<String> newRightFieldNames = new ArrayList<>();
    int leftKeyCount = leftJoinKeys.size();
    int i;
    for (i = 0; i < origLeftInputSize; i++) {
        final RelDataTypeField field = leftRel.getRowType().getFieldList().get(i);
        newLeftFields.add(rexBuilder.makeInputRef(field.getType(), i));
        newLeftFieldNames.add(field.getName());
    }
    for (i = 0; i < origRightInputSize; i++) {
        final RelDataTypeField field = rightRel.getRowType().getFieldList().get(i);
        newRightFields.add(rexBuilder.makeInputRef(field.getType(), i));
        newRightFieldNames.add(field.getName());
    }
    ImmutableBitSet.Builder origColEqCondsPosBuilder = ImmutableBitSet.builder();
    int newKeyCount = 0;
    List<Pair<Integer, Integer>> origColEqConds = new ArrayList<>();
    for (i = 0; i < leftKeyCount; i++) {
        RexNode leftKey = leftJoinKeys.get(i);
        RexNode rightKey = rightJoinKeys.get(i);
        if (leftKey instanceof RexInputRef && rightKey instanceof RexInputRef) {
            origColEqConds.add(Pair.of(((RexInputRef) leftKey).getIndex(), ((RexInputRef) rightKey).getIndex()));
            origColEqCondsPosBuilder.set(i);
        } else {
            newLeftFields.add(leftKey);
            newLeftFieldNames.add(null);
            newRightFields.add(rightKey);
            newRightFieldNames.add(null);
            newKeyCount++;
        }
    }
    ImmutableBitSet origColEqCondsPos = origColEqCondsPosBuilder.build();
    for (i = 0; i < origColEqConds.size(); i++) {
        Pair<Integer, Integer> p = origColEqConds.get(i);
        int condPos = origColEqCondsPos.nth(i);
        RexNode leftKey = leftJoinKeys.get(condPos);
        RexNode rightKey = rightJoinKeys.get(condPos);
        leftKeys.add(p.left);
        rightKeys.add(p.right);
        RexNode cond = rexBuilder.makeCall(SqlStdOperatorTable.EQUALS, rexBuilder.makeInputRef(leftKey.getType(), systemColCount + p.left), rexBuilder.makeInputRef(rightKey.getType(), systemColCount + origLeftInputSize + newKeyCount + p.right));
        if (outJoinCond == null) {
            outJoinCond = cond;
        } else {
            outJoinCond = rexBuilder.makeCall(SqlStdOperatorTable.AND, outJoinCond, cond);
        }
    }
    if (newKeyCount == 0) {
        return outJoinCond;
    }
    int newLeftOffset = systemColCount + origLeftInputSize;
    int newRightOffset = systemColCount + origLeftInputSize + origRightInputSize + newKeyCount;
    for (i = 0; i < newKeyCount; i++) {
        leftKeys.add(origLeftInputSize + i);
        rightKeys.add(origRightInputSize + i);
        RexNode cond = rexBuilder.makeCall(SqlStdOperatorTable.EQUALS, rexBuilder.makeInputRef(newLeftFields.get(origLeftInputSize + i).getType(), newLeftOffset + i), rexBuilder.makeInputRef(newRightFields.get(origRightInputSize + i).getType(), newRightOffset + i));
        if (outJoinCond == null) {
            outJoinCond = cond;
        } else {
            outJoinCond = rexBuilder.makeCall(SqlStdOperatorTable.AND, outJoinCond, cond);
        }
    }
    // added project if need to produce new keys than the original input fields
    if (newKeyCount > 0) {
        leftRel = factory.createProject(leftRel, Collections.emptyList(), newLeftFields, SqlValidatorUtil.uniquify(newLeftFieldNames, false));
        rightRel = factory.createProject(rightRel, Collections.emptyList(), newRightFields, SqlValidatorUtil.uniquify(newRightFieldNames, false));
    }
    inputRels[0] = leftRel;
    inputRels[1] = rightRel;
    return outJoinCond;
}

Example 80 with Pair

use of org.apache.calcite.util.Pair in project flink by apache.

the class HiveParserCalcitePlanner method getWindowRexAndType.

private Pair<RexNode, TypeInfo> getWindowRexAndType(HiveParserWindowingSpec.WindowExpressionSpec winExprSpec, RelNode srcRel) throws SemanticException {
    RexNode window;
    if (winExprSpec instanceof HiveParserWindowingSpec.WindowFunctionSpec) {
        HiveParserWindowingSpec.WindowFunctionSpec wFnSpec = (HiveParserWindowingSpec.WindowFunctionSpec) winExprSpec;
        HiveParserASTNode windowProjAst = wFnSpec.getExpression();
        // TODO: do we need to get to child?
        int wndSpecASTIndx = getWindowSpecIndx(windowProjAst);
        // 2. Get Hive Aggregate Info
        AggInfo hiveAggInfo = getHiveAggInfo(windowProjAst, wndSpecASTIndx - 1, relToRowResolver.get(srcRel), (HiveParserWindowingSpec.WindowFunctionSpec) winExprSpec, semanticAnalyzer, frameworkConfig, cluster);
        // 3. Get Calcite Return type for Agg Fn
        RelDataType calciteAggFnRetType = HiveParserUtils.toRelDataType(hiveAggInfo.getReturnType(), cluster.getTypeFactory());
        // 4. Convert Agg Fn args to Calcite
        Map<String, Integer> posMap = relToHiveColNameCalcitePosMap.get(srcRel);
        HiveParserRexNodeConverter converter = new HiveParserRexNodeConverter(cluster, srcRel.getRowType(), posMap, 0, false, funcConverter);
        List<RexNode> calciteAggFnArgs = new ArrayList<>();
        List<RelDataType> calciteAggFnArgTypes = new ArrayList<>();
        for (int i = 0; i < hiveAggInfo.getAggParams().size(); i++) {
            calciteAggFnArgs.add(converter.convert(hiveAggInfo.getAggParams().get(i)));
            calciteAggFnArgTypes.add(HiveParserUtils.toRelDataType(hiveAggInfo.getAggParams().get(i).getTypeInfo(), cluster.getTypeFactory()));
        }
        // 5. Get Calcite Agg Fn
        final SqlAggFunction calciteAggFn = HiveParserSqlFunctionConverter.getCalciteAggFn(hiveAggInfo.getUdfName(), hiveAggInfo.isDistinct(), calciteAggFnArgTypes, calciteAggFnRetType);
        // 6. Translate Window spec
        HiveParserRowResolver inputRR = relToRowResolver.get(srcRel);
        HiveParserWindowingSpec.WindowSpec wndSpec = ((HiveParserWindowingSpec.WindowFunctionSpec) winExprSpec).getWindowSpec();
        List<RexNode> partitionKeys = getPartitionKeys(wndSpec.getPartition(), converter, inputRR, new HiveParserTypeCheckCtx(inputRR, frameworkConfig, cluster), semanticAnalyzer);
        List<RexFieldCollation> orderKeys = getOrderKeys(wndSpec.getOrder(), converter, inputRR, new HiveParserTypeCheckCtx(inputRR, frameworkConfig, cluster), semanticAnalyzer);
        RexWindowBound lowerBound = getBound(wndSpec.getWindowFrame().getStart(), cluster);
        RexWindowBound upperBound = getBound(wndSpec.getWindowFrame().getEnd(), cluster);
        boolean isRows = wndSpec.getWindowFrame().getWindowType() == HiveParserWindowingSpec.WindowType.ROWS;
        window = HiveParserUtils.makeOver(cluster.getRexBuilder(), calciteAggFnRetType, calciteAggFn, calciteAggFnArgs, partitionKeys, orderKeys, lowerBound, upperBound, isRows, true, false, false, false);
        window = window.accept(funcConverter);
    } else {
        throw new SemanticException("Unsupported window Spec");
    }
    return new Pair<>(window, HiveParserTypeConverter.convert(window.getType()));
}