Java Code Examples for org.apache.samza.operators.KV#of()

public KV<K, V> fromBytes(byte[] bytes) {
  if (bytes != null) {
    ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
    int keyLength = byteBuffer.getInt();
    byte[] keyBytes = new byte[keyLength];
    byteBuffer.get(keyBytes);
    int valueLength = byteBuffer.getInt();
    byte[] valueBytes = new byte[valueLength];
    byteBuffer.get(valueBytes);
    K key = keySerde.fromBytes(keyBytes);
    V value = valueSerde.fromBytes(valueBytes);
    return KV.of(key, value);
  } else {
    return null;
  }
}
 

@Override
protected CompletionStage<Collection<Object>> handleMessageAsync(IncomingMessageEnvelope message,
    MessageCollector collector, TaskCoordinator coordinator) {
  Object result;
  InputTransformer transformer = inputOpSpec.getTransformer();
  if (transformer != null) {
    result = transformer.apply(message);
  } else {
    result = this.inputOpSpec.isKeyed() ? KV.of(message.getKey(), message.getMessage()) : message.getMessage();
  }

  Collection<Object> output = Optional.ofNullable(result)
      .map(Collections::singletonList)
      .orElse(Collections.emptyList());

  return CompletableFuture.completedFuture(output);
}
 

private Collection<JM> getJoinOutput(K key, Object value, M message) {
  R record = value == null ? null : (R) KV.of(key, value);

  JM output = joinOpSpec.getJoinFn().apply(message, record);

  // The support for inner and outer join will be provided in the jonFn. For inner join, the joinFn might
  // return null, when the corresponding record is absent in the table.
  return output != null ?
      Collections.singletonList(output) : Collections.emptyList();
}
 

@Override
public Object apply(IncomingMessageEnvelope ime) {
  Assert.notNull(ime, "ime is null");
  KV<Object, Object> keyAndMessageKV = KV.of(ime.getKey(), ime.getMessage());
  SamzaSqlRelMsgMetadata metadata = new SamzaSqlRelMsgMetadata(ime.getEventTime(), ime.getArrivalTime());
  SamzaSqlInputMessage samzaMsg = SamzaSqlInputMessage.of(keyAndMessageKV, metadata);
  return  samzaMsg;
}
 

@Test
public void testWithInnerJoinWithTableOnRight() {
  SamzaSqlRelMessage streamMsg = new SamzaSqlRelMessage(streamFieldNames, streamFieldValues, new SamzaSqlRelMsgMetadata(0L, 0L));
  SamzaSqlRelMessage tableMsg = new SamzaSqlRelMessage(tableFieldNames, tableFieldValues, new SamzaSqlRelMsgMetadata(0L, 0L));
  JoinRelType joinRelType = JoinRelType.INNER;
  List<Integer> streamKeyIds = Arrays.asList(0, 1);
  List<Integer> tableKeyIds = Arrays.asList(0, 1);
  SamzaSqlRelRecord compositeKey = SamzaSqlRelMessage.createSamzaSqlCompositeKey(tableMsg, tableKeyIds);
  KV<SamzaSqlRelRecord, SamzaSqlRelMessage> record = KV.of(compositeKey, tableMsg);

  JoinInputNode mockTableInputNode = mock(JoinInputNode.class);
  when(mockTableInputNode.getKeyIds()).thenReturn(tableKeyIds);
  when(mockTableInputNode.isPosOnRight()).thenReturn(true);
  when(mockTableInputNode.getFieldNames()).thenReturn(tableFieldNames);

  JoinInputNode mockStreamInputNode = mock(JoinInputNode.class);
  when(mockStreamInputNode.getKeyIds()).thenReturn(streamKeyIds);
  when(mockStreamInputNode.isPosOnRight()).thenReturn(false);
  when(mockStreamInputNode.getFieldNames()).thenReturn(streamFieldNames);

  SamzaSqlLocalTableJoinFunction joinFn =
      new SamzaSqlLocalTableJoinFunction(mockStreamInputNode, mockTableInputNode, joinRelType);
  SamzaSqlRelMessage outMsg = joinFn.apply(streamMsg, record);

  Assert.assertEquals(outMsg.getSamzaSqlRelRecord().getFieldValues().size(),
      outMsg.getSamzaSqlRelRecord().getFieldNames().size());
  List<String> expectedFieldNames = new ArrayList<>(streamFieldNames);
  expectedFieldNames.addAll(tableFieldNames);
  List<Object> expectedFieldValues = new ArrayList<>(streamFieldValues);
  expectedFieldValues.addAll(tableFieldValues);
  Assert.assertEquals(outMsg.getSamzaSqlRelRecord().getFieldValues(), expectedFieldValues);
}
 

@Test
public void testWithInnerJoinWithTableOnLeft() {
  SamzaSqlRelMessage streamMsg = new SamzaSqlRelMessage(streamFieldNames, streamFieldValues, new SamzaSqlRelMsgMetadata(0L, 0L));
  SamzaSqlRelMessage tableMsg = new SamzaSqlRelMessage(tableFieldNames, tableFieldValues, new SamzaSqlRelMsgMetadata(0L, 0L));
  JoinRelType joinRelType = JoinRelType.INNER;
  List<Integer> streamKeyIds = Arrays.asList(0, 2);
  List<Integer> tableKeyIds = Arrays.asList(0, 2);
  SamzaSqlRelRecord compositeKey = SamzaSqlRelMessage.createSamzaSqlCompositeKey(tableMsg, tableKeyIds);
  KV<SamzaSqlRelRecord, SamzaSqlRelMessage> record = KV.of(compositeKey, tableMsg);

  JoinInputNode mockTableInputNode = mock(JoinInputNode.class);
  when(mockTableInputNode.getKeyIds()).thenReturn(tableKeyIds);
  when(mockTableInputNode.isPosOnRight()).thenReturn(false);
  when(mockTableInputNode.getFieldNames()).thenReturn(tableFieldNames);

  JoinInputNode mockStreamInputNode = mock(JoinInputNode.class);
  when(mockStreamInputNode.getKeyIds()).thenReturn(streamKeyIds);
  when(mockStreamInputNode.isPosOnRight()).thenReturn(true);
  when(mockStreamInputNode.getFieldNames()).thenReturn(streamFieldNames);

  SamzaSqlLocalTableJoinFunction joinFn =
      new SamzaSqlLocalTableJoinFunction(mockStreamInputNode, mockTableInputNode, joinRelType);
  SamzaSqlRelMessage outMsg = joinFn.apply(streamMsg, record);

  Assert.assertEquals(outMsg.getSamzaSqlRelRecord().getFieldValues().size(),
      outMsg.getSamzaSqlRelRecord().getFieldNames().size());
  List<String> expectedFieldNames = new ArrayList<>(tableFieldNames);
  expectedFieldNames.addAll(streamFieldNames);
  List<Object> expectedFieldValues = new ArrayList<>(tableFieldValues);
  expectedFieldValues.addAll(streamFieldValues);
  Assert.assertEquals(outMsg.getSamzaSqlRelRecord().getFieldValues(), expectedFieldValues);
}
 

@Test
public void testJoinChain() {
  String inputStreamId1 = "input1";
  String inputStreamId2 = "input2";
  String inputSystem = "input-system";
  String inputPhysicalName1 = "input-stream1";
  String inputPhysicalName2 = "input-stream2";
  HashMap<String, String> configs = new HashMap<>();
  configs.put(JobConfig.JOB_NAME, "jobName");
  configs.put(JobConfig.JOB_ID, "jobId");
  StreamTestUtils.addStreamConfigs(configs, inputStreamId1, inputSystem, inputPhysicalName1);
  StreamTestUtils.addStreamConfigs(configs, inputStreamId2, inputSystem, inputPhysicalName2);
  Config config = new MapConfig(configs);
  when(this.context.getJobContext().getConfig()).thenReturn(config);

  Integer joinKey = new Integer(1);
  Function<Object, Integer> keyFn = (Function & Serializable) m -> joinKey;
  JoinFunction testJoinFunction = new TestJoinFunction("jobName-jobId-join-j1",
      (BiFunction & Serializable) (m1, m2) -> KV.of(m1, m2), keyFn, keyFn);

  StreamApplicationDescriptorImpl graphSpec = new StreamApplicationDescriptorImpl(appDesc -> {
    GenericSystemDescriptor sd = new GenericSystemDescriptor(inputSystem, "mockFactoryClass");
    GenericInputDescriptor inputDescriptor1 = sd.getInputDescriptor(inputStreamId1, mock(Serde.class));
    GenericInputDescriptor inputDescriptor2 = sd.getInputDescriptor(inputStreamId2, mock(Serde.class));
    MessageStream<Object> inputStream1 = appDesc.getInputStream(inputDescriptor1);
    MessageStream<Object> inputStream2 = appDesc.getInputStream(inputDescriptor2);

    inputStream1.join(inputStream2, testJoinFunction,
        mock(Serde.class), mock(Serde.class), mock(Serde.class), Duration.ofHours(1), "j1");
  }, config);

  TaskName mockTaskName = mock(TaskName.class);
  TaskModel taskModel = mock(TaskModel.class);
  when(taskModel.getTaskName()).thenReturn(mockTaskName);
  when(this.context.getTaskContext().getTaskModel()).thenReturn(taskModel);

  KeyValueStore mockLeftStore = mock(KeyValueStore.class);
  when(this.context.getTaskContext().getStore(eq("jobName-jobId-join-j1-L"))).thenReturn(mockLeftStore);
  KeyValueStore mockRightStore = mock(KeyValueStore.class);
  when(this.context.getTaskContext().getStore(eq("jobName-jobId-join-j1-R"))).thenReturn(mockRightStore);
  OperatorImplGraph opImplGraph =
      new OperatorImplGraph(graphSpec.getOperatorSpecGraph(), this.context, mock(Clock.class));

  // verify that join function is initialized once.
  assertEquals(TestJoinFunction.getInstanceByTaskName(mockTaskName, "jobName-jobId-join-j1").numInitCalled, 1);

  InputOperatorImpl inputOpImpl1 = opImplGraph.getInputOperator(new SystemStream(inputSystem, inputPhysicalName1));
  InputOperatorImpl inputOpImpl2 = opImplGraph.getInputOperator(new SystemStream(inputSystem, inputPhysicalName2));
  PartialJoinOperatorImpl leftPartialJoinOpImpl =
      (PartialJoinOperatorImpl) inputOpImpl1.registeredOperators.iterator().next();
  PartialJoinOperatorImpl rightPartialJoinOpImpl =
      (PartialJoinOperatorImpl) inputOpImpl2.registeredOperators.iterator().next();

  assertEquals(leftPartialJoinOpImpl.getOperatorSpec(), rightPartialJoinOpImpl.getOperatorSpec());
  assertNotSame(leftPartialJoinOpImpl, rightPartialJoinOpImpl);

  // verify that left partial join operator calls getFirstKey
  Object mockLeftMessage = mock(Object.class);
  long currentTimeMillis = System.currentTimeMillis();
  when(mockLeftStore.get(eq(joinKey))).thenReturn(new TimestampedValue<>(mockLeftMessage, currentTimeMillis));
  IncomingMessageEnvelope leftMessage = new IncomingMessageEnvelope(mock(SystemStreamPartition.class), "", "", mockLeftMessage);
  inputOpImpl1.onMessage(leftMessage, mock(MessageCollector.class), mock(TaskCoordinator.class));

  // verify that right partial join operator calls getSecondKey
  Object mockRightMessage = mock(Object.class);
  when(mockRightStore.get(eq(joinKey))).thenReturn(new TimestampedValue<>(mockRightMessage, currentTimeMillis));
  IncomingMessageEnvelope rightMessage = new IncomingMessageEnvelope(mock(SystemStreamPartition.class), "", "", mockRightMessage);
  inputOpImpl2.onMessage(rightMessage, mock(MessageCollector.class), mock(TaskCoordinator.class));


  // verify that the join function apply is called with the correct messages on match
  assertEquals(((TestJoinFunction) TestJoinFunction.getInstanceByTaskName(mockTaskName, "jobName-jobId-join-j1")).joinResults.size(), 1);
  KV joinResult = (KV) ((TestJoinFunction) TestJoinFunction.getInstanceByTaskName(mockTaskName, "jobName-jobId-join-j1")).joinResults.iterator().next();
  assertEquals(joinResult.getKey(), mockLeftMessage);
  assertEquals(joinResult.getValue(), mockRightMessage);
}
 

@Test
public void testWithInnerJoinWithTableOnRight() {
  Map<String, String> props = new HashMap<>();
  SystemStream ss = new SystemStream("test", "nestedRecord");
  props.put(
      String.format(ConfigBasedAvroRelSchemaProviderFactory.CFG_SOURCE_SCHEMA, ss.getSystem(), ss.getStream()),
      SimpleRecord.SCHEMA$.toString());
  ConfigBasedAvroRelSchemaProviderFactory factory = new ConfigBasedAvroRelSchemaProviderFactory();
  AvroRelSchemaProvider schemaProvider =
      (AvroRelSchemaProvider) factory.create(ss, new MapConfig(props));
  AvroRelConverter relConverter =
      new AvroRelConverter(ss, schemaProvider, new MapConfig());
  SamzaRelTableKeyConverter relTableKeyConverter = new SampleRelTableKeyConverter();
  String remoteTableName = "testDb.testTable.$table";

  GenericData.Record tableRecord = new GenericData.Record(SimpleRecord.SCHEMA$);
  tableRecord.put("id", 1);
  tableRecord.put("name", "name1");

  SamzaSqlRelMessage streamMsg = new SamzaSqlRelMessage(streamFieldNames, streamFieldValues,
      new SamzaSqlRelMsgMetadata(0L, 0L));
  SamzaSqlRelMessage tableMsg = relConverter.convertToRelMessage(new KV(tableRecord.get("id"), tableRecord));
  JoinRelType joinRelType = JoinRelType.INNER;
  List<Integer> streamKeyIds = Arrays.asList(1);
  List<Integer> tableKeyIds = Arrays.asList(0);
  KV<Object, GenericRecord> record = KV.of(tableRecord.get("id"), tableRecord);

  JoinInputNode mockTableInputNode = mock(JoinInputNode.class);
  when(mockTableInputNode.getKeyIds()).thenReturn(tableKeyIds);
  when(mockTableInputNode.isPosOnRight()).thenReturn(true);
  when(mockTableInputNode.getFieldNames()).thenReturn(tableMsg.getSamzaSqlRelRecord().getFieldNames());
  when(mockTableInputNode.getSourceName()).thenReturn(remoteTableName);

  JoinInputNode mockStreamInputNode = mock(JoinInputNode.class);
  when(mockStreamInputNode.getKeyIds()).thenReturn(streamKeyIds);
  when(mockStreamInputNode.isPosOnRight()).thenReturn(false);
  when(mockStreamInputNode.getFieldNames()).thenReturn(streamFieldNames);

  SamzaSqlRemoteTableJoinFunction joinFn =
      new SamzaSqlRemoteTableJoinFunction(relConverter, relTableKeyConverter, mockStreamInputNode, mockTableInputNode,
          joinRelType, 0);
  SamzaSqlRelMessage outMsg = joinFn.apply(streamMsg, record);

  Assert.assertEquals(outMsg.getSamzaSqlRelRecord().getFieldValues().size(),
      outMsg.getSamzaSqlRelRecord().getFieldNames().size());
  List<String> expectedFieldNames = new ArrayList<>(streamFieldNames);
  expectedFieldNames.addAll(tableMsg.getSamzaSqlRelRecord().getFieldNames());
  List<Object> expectedFieldValues = new ArrayList<>(streamFieldValues);
  expectedFieldValues.addAll(tableMsg.getSamzaSqlRelRecord().getFieldValues());

  Assert.assertEquals(expectedFieldNames, outMsg.getSamzaSqlRelRecord().getFieldNames());
  Assert.assertEquals(expectedFieldValues, outMsg.getSamzaSqlRelRecord().getFieldValues());
}
 

@Override
public KV<String, JsonObject> apply(KV<String, String> message, KV<String, String> record) {
  return KV.of(message.getKey(), JsonObject.create().put("name", message.key).put("age", record.getValue()));
}