Examples with DataContext org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext used on opensource projects

Example 11 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project samza by apache.

the class TestFilterTranslator method testTranslate.

@Test
public void testTranslate() throws IOException, ClassNotFoundException {
    // setup mock values to the constructor of FilterTranslator
    LogicalFilter mockFilter = PowerMockito.mock(LogicalFilter.class);
    Context mockContext = mock(Context.class);
    ContainerContext mockContainerContext = mock(ContainerContext.class);
    TranslatorContext mockTranslatorContext = mock(TranslatorContext.class);
    TestMetricsRegistryImpl metricsRegistry = new TestMetricsRegistryImpl();
    RelNode mockInput = mock(RelNode.class);
    when(mockFilter.getInput()).thenReturn(mockInput);
    when(mockInput.getId()).thenReturn(1);
    when(mockFilter.getId()).thenReturn(2);
    StreamApplicationDescriptorImpl mockGraph = mock(StreamApplicationDescriptorImpl.class);
    OperatorSpec<Object, SamzaSqlRelMessage> mockInputOp = mock(OperatorSpec.class);
    MessageStream<SamzaSqlRelMessage> mockStream = new MessageStreamImpl<>(mockGraph, 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(metricsRegistry);
    // Apply translate() method to verify that we are getting the correct filter operator constructed
    FilterTranslator filterTranslator = new FilterTranslator(1);
    filterTranslator.translate(mockFilter, LOGICAL_OP_ID, mockTranslatorContext);
    // make sure that context has been registered with LogicFilter and output message streams
    verify(mockTranslatorContext, times(1)).registerRelNode(2, mockFilter);
    verify(mockTranslatorContext, times(1)).registerMessageStream(2, this.getRegisteredMessageStream(2));
    when(mockTranslatorContext.getRelNode(2)).thenReturn(mockFilter);
    when(mockTranslatorContext.getMessageStream(2)).thenReturn(this.getRegisteredMessageStream(2));
    StreamOperatorSpec filterSpec = (StreamOperatorSpec) Whitebox.getInternalState(this.getRegisteredMessageStream(2), "operatorSpec");
    assertNotNull(filterSpec);
    assertEquals(filterSpec.getOpCode(), OperatorSpec.OpCode.FILTER);
    // Verify that the describe() method will establish the context for the filter function
    Map<Integer, TranslatorContext> mockContexts = new HashMap<>();
    mockContexts.put(1, mockTranslatorContext);
    when(mockContext.getApplicationTaskContext()).thenReturn(new SamzaSqlApplicationContext(mockContexts));
    filterSpec.getTransformFn().init(mockContext);
    FilterFunction filterFn = (FilterFunction) Whitebox.getInternalState(filterSpec, "filterFn");
    assertNotNull(filterFn);
    assertEquals(mockTranslatorContext, Whitebox.getInternalState(filterFn, "translatorContext"));
    assertEquals(mockFilter, Whitebox.getInternalState(filterFn, "filter"));
    assertEquals(mockExpr, Whitebox.getInternalState(filterFn, "expr"));
    // Verify MetricsRegistry works with Project
    assertEquals(1, metricsRegistry.getGauges().size());
    assertTrue(metricsRegistry.getGauges().get(LOGICAL_OP_ID).size() > 0);
    assertEquals(1, metricsRegistry.getCounters().size());
    assertEquals(3, metricsRegistry.getCounters().get(LOGICAL_OP_ID).size());
    assertEquals(0, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(0, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
    // Calling filterFn.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];
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = new Boolean(true);
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    assertTrue(filterFn.apply(mockInputMsg));
    doAnswer(invocation -> {
        Object[] retValue = invocation.getArgumentAt(4, Object[].class);
        retValue[0] = new Boolean(false);
        return null;
    }).when(mockExpr).execute(eq(executionContext), eq(mockContext), eq(dataContext), eq(mockInputMsg.getSamzaSqlRelRecord().getFieldValues().toArray()), eq(result));
    assertFalse(filterFn.apply(mockInputMsg));
    // Verify filterFn.apply() updates the MetricsRegistry metrics
    assertEquals(2, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(0).getCount());
    assertEquals(1, metricsRegistry.getCounters().get(LOGICAL_OP_ID).get(1).getCount());
}

Example 12 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project storm by apache.

the class TestPlanCompiler method testDateKeywordsAndFunctions.

/**
 * All the date/time/timestamp related tests are done here, because Avatica converts the result of date functions to java.sql classes
 * whereas Stream provides long type which makes different semantic from Stream implementation.
 */
@Test
public void testDateKeywordsAndFunctions() throws Exception {
    int EXPECTED_VALUE_SIZE = 1;
    String sql = "SELECT " + "LOCALTIME, CURRENT_TIME, LOCALTIMESTAMP, CURRENT_TIMESTAMP, CURRENT_DATE, " + "DATE '1970-05-15' AS datefield, TIME '00:00:00' AS timefield, TIMESTAMP '2016-01-01 00:00:00' as timestampfield, " + "EXTRACT(MONTH FROM TIMESTAMP '2010-01-23 12:34:56')," + "FLOOR(DATE '2016-01-23' TO MONTH)," + "CEIL(TIME '12:34:56' TO MINUTE)," + "{fn CURDATE()} = CURRENT_DATE, {fn CURTIME()} = LOCALTIME, {fn NOW()} = LOCALTIMESTAMP," + "{fn QUARTER(DATE '2016-10-07')}, {fn TIMESTAMPADD(MINUTE, 15, TIMESTAMP '2016-10-07 00:00:00')}," + "{fn TIMESTAMPDIFF(SECOND, TIMESTAMP '2016-10-06 00:00:00', TIMESTAMP '2016-10-07 00:00:00')}," + "INTERVAL '1-5' YEAR TO MONTH AS intervalfield, " + "(DATE '1970-01-01', DATE '1970-01-15') AS anchoredinterval_field " + "FROM FOO " + "WHERE ID > 0 AND ID < 2";
    TestCompilerUtils.CalciteState state = TestCompilerUtils.sqlOverDummyTable(sql);
    final Map<String, ISqlStreamsDataSource> data = new HashMap<>();
    data.put("FOO", new TestUtils.MockSqlStreamsDataSource());
    QueryPlanner planner = new QueryPlanner(state.schema());
    AbstractStreamsProcessor proc = planner.compile(data, sql);
    // inject output bolt
    proc.outputStream().to(new TestUtils.MockBolt());
    final DataContext dataContext = proc.getDataContext();
    final StormTopology topo = proc.build();
    SqlTestUtil.runStormTopology(cluster, TestUtils.MockBolt.getCollectedValues(), EXPECTED_VALUE_SIZE, proc, topo);
    long utcTimestamp = (long) dataContext.get(DataContext.Variable.UTC_TIMESTAMP.camelName);
    long currentTimestamp = (long) dataContext.get(DataContext.Variable.CURRENT_TIMESTAMP.camelName);
    long localTimestamp = (long) dataContext.get(DataContext.Variable.LOCAL_TIMESTAMP.camelName);
    System.out.println(TestUtils.MockBolt.getCollectedValues());
    java.sql.Timestamp timestamp = new java.sql.Timestamp(utcTimestamp);
    int dateInt = (int) timestamp.toLocalDateTime().atOffset(ZoneOffset.UTC).toLocalDate().toEpochDay();
    int localTimeInt = (int) (localTimestamp % DateTimeUtils.MILLIS_PER_DAY);
    int currentTimeInt = (int) (currentTimestamp % DateTimeUtils.MILLIS_PER_DAY);
    Assert.assertArrayEquals(new Values[] { new Values(localTimeInt, currentTimeInt, localTimestamp, currentTimestamp, dateInt, 134, 0, 1451606400000L, 1L, 0L, 45300000, true, true, true, 4L, 1475799300000L, 86400, 17, 0, 14) }, TestUtils.MockBolt.getCollectedValues().toArray());
}

Example 13 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project storm by apache.

the class QueryPlanner method compile.

public AbstractStreamsProcessor compile(Map<String, ISqlStreamsDataSource> sources, String query) throws Exception {
    StreamsRel relNode = getPlan(query);
    StreamsPlanCreator streamsPlanCreator = new StreamsPlanCreator(sources, new RexBuilder(typeFactory));
    relNode.streamsPlan(streamsPlanCreator);
    final StreamBuilder streamBuilder = streamsPlanCreator.getStreamBuilder();
    final Stream<Values> lastStream = streamsPlanCreator.pop();
    final DataContext dc = streamsPlanCreator.getDataContext();
    final List<CompilingClassLoader> cls = streamsPlanCreator.getClassLoaders();
    return new AbstractStreamsProcessor() {

        @Override
        public StormTopology build() {
            return streamBuilder.build();
        }

        @Override
        public Stream<Values> outputStream() {
            return lastStream;
        }

        @Override
        public DataContext getDataContext() {
            return dc;
        }

        @Override
        public List<CompilingClassLoader> getClassLoaders() {
            return cls;
        }
    };
}

Example 14 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project storm by apache.

the class EvaluationCalc method apply.

@Override
public Iterable<Values> apply(Values input) {
    Context calciteContext = new StormContext(dataContext);
    calciteContext.values = input.toArray();
    if (filterInstance != null) {
        filterInstance.execute(calciteContext, outputValues);
        // filtered out
        if (outputValues[0] == null || !((Boolean) outputValues[0])) {
            return Collections.emptyList();
        }
    }
    if (projectionInstance != null) {
        projectionInstance.execute(calciteContext, outputValues);
        return Collections.singletonList(new Values(outputValues));
    } else {
        return Collections.singletonList(new Values(input.toArray()));
    }
}

Example 15 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project storm by apache.

the class EvaluationFilter method test.

@Override
public boolean test(Values input) {
    Context calciteContext = new StormContext(dataContext);
    calciteContext.values = input.toArray();
    filterInstance.execute(calciteContext, outputValues);
    return (outputValues[0] != null && (boolean) outputValues[0]);
}