ProjectTranslator { void translate(final Project project, final String logicalOpId, final TranslatorContext context) { MessageStream<SamzaSqlRelMessage> messageStream = context.getMessageStream(project.getInput().getId()); final int projectId = project.getId(); MessageStream<SamzaSqlRelMessage> outputStream = messageStream.map(new ProjectMapFunction(projectId, queryId, logicalOpId)); List<RexNode> projects = project.getProjects(); List<Integer> flattenProjects = IntStream.range(0, projects.size()) .filter(i -> this.isFlatten(projects.get(i))) .boxed() .collect(Collectors.toList()); if (flattenProjects.size() > 0) { if (flattenProjects.size() > 1) { String msg = "Multiple flatten operators in a single query is not supported"; LOG.error(msg); throw new SamzaException(msg); } outputStream = translateFlatten(flattenProjects.get(0), outputStream); } context.registerMessageStream(project.getId(), outputStream); context.registerRelNode(project.getId(), project); } ProjectTranslator(int queryId); static SamzaSqlRelRecord buildSamzaRelRecord(Object[] objects, RelDataType rowType); static Object[] convertToJavaRow(SamzaSqlRelRecord samzaSqlRelRecord);  }
@Test public void testTranslate() throws IOException, ClassNotFoundException { 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); List<RelDataTypeField> relFields = new ArrayList<>(); String fieldName = TEST_FIELD; int fieldPos = 0; RelDataType dataType = mock(RelDataType.class); when(dataType.getSqlTypeName()).thenReturn(SqlTypeName.ANY); relFields.add(new RelDataTypeFieldImpl(fieldName, fieldPos, dataType)); when(mockRowType.getFieldCount()).thenReturn(1); when(mockProject.getRowType()).thenReturn(mockRowType); when(mockProject.getRowType().getSqlTypeName()).thenReturn(SqlTypeName.ROW); when(mockProject.getRowType().getFieldList()).thenReturn(relFields); when(mockProject.getRowType().isStruct()).thenReturn(true); RexNode mockRexField = mock(RexNode.class); List<Pair<RexNode, String>> namedProjects = new ArrayList<>(); namedProjects.add(Pair.of(mockRexField, TEST_FIELD)); when(mockProject.getNamedProjects()).thenReturn(namedProjects); when(mockProject.getRowType()).thenReturn(mockRowType); when(mockProject.getRowType().getFieldNames()).thenReturn(ImmutableList.of(TEST_FIELD)); 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); ProjectTranslator projectTranslator = new ProjectTranslator(1); projectTranslator.translate(mockProject, LOGICAL_OP_ID, mockTranslatorContext); 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); 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")); 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()); 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)); assertEquals(1, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(0).getCount()); assertEquals(1, testMetricsRegistryImpl.getCounters().get(LOGICAL_OP_ID).get(1).getCount()); }