| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| */ |
| |
| package org.apache.samza.sql.system; |
| |
| import java.nio.ByteBuffer; |
| import java.util.ArrayList; |
| import java.util.HashMap; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| import java.util.stream.Collectors; |
| import java.util.stream.IntStream; |
| import org.apache.avro.generic.GenericData; |
| import org.apache.avro.generic.GenericRecord; |
| import org.apache.samza.config.Config; |
| import org.apache.samza.metrics.MetricsRegistry; |
| import org.apache.samza.sql.avro.schemas.AddressRecord; |
| import org.apache.samza.sql.avro.schemas.Company; |
| import org.apache.samza.sql.avro.schemas.ComplexRecord; |
| import org.apache.samza.sql.avro.schemas.Kind; |
| import org.apache.samza.sql.avro.schemas.MyFixed; |
| import org.apache.samza.sql.avro.schemas.PageView; |
| import org.apache.samza.sql.avro.schemas.PhoneNumber; |
| import org.apache.samza.sql.avro.schemas.Profile; |
| import org.apache.samza.sql.avro.schemas.SimpleRecord; |
| import org.apache.samza.sql.avro.schemas.StreetNumRecord; |
| import org.apache.samza.sql.avro.schemas.SubRecord; |
| import org.apache.samza.system.IncomingMessageEnvelope; |
| import org.apache.samza.system.OutgoingMessageEnvelope; |
| import org.apache.samza.system.SystemAdmin; |
| import org.apache.samza.system.SystemConsumer; |
| import org.apache.samza.system.SystemFactory; |
| import org.apache.samza.system.SystemProducer; |
| import org.apache.samza.system.SystemStreamPartition; |
| import org.junit.Assert; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| |
| public class TestAvroSystemFactory implements SystemFactory { |
| private static final Logger LOG = LoggerFactory.getLogger(TestAvroSystemFactory.class); |
| |
| public static final String CFG_NUM_MESSAGES = "numMessages"; |
| public static final String CFG_INCLUDE_NULL_FOREIGN_KEYS = "includeNullForeignKeys"; |
| public static final String CFG_INCLUDE_NULL_SIMPLE_RECORDS = "includeNullSimpleRecords"; |
| public static final String CFG_SLEEP_BETWEEN_POLLS_MS = "sleepBetweenPollsMs"; |
| |
| private static final String[] PROFILE_NAMES = {"John", "Mike", "Mary", "Joe", "Brad", "Jennifer"}; |
| private static final int[] PROFILE_ZIPS = {94000, 94001, 94002, 94003, 94004, 94005}; |
| private static final int[] STREET_NUMS = {1234, 1235, 1236, 1237, 1238, 1239}; |
| private static final String[] PHONE_NUMBERS = {"000-000-0000", "111-111-1111", "222-222-2222", "333-333-3333", |
| "444-444-4444", "555-555-5555"}; |
| public static final String[] COMPANIES = {"MSFT", "LKND", "GOOG", "FB", "AMZN", "CSCO"}; |
| public static final String[] PAGE_KEYS = {"inbox", "home", "search", "pymk", "group", "job"}; |
| public static final byte[] DEFAULT_TRACKING_ID_BYTES = |
| {76, 75, -24, 10, 33, -117, 24, -52, -110, -39, -5, 102, 65, 57, -62, -1}; |
| public static final int NULL_RECORD_FREQUENCY = 5; |
| |
| |
| public static List<OutgoingMessageEnvelope> messages = new ArrayList<>(); |
| |
| public static List<String> getPageKeyProfileNameJoin(int numMessages) { |
| return IntStream.range(0, numMessages) |
| .mapToObj(i -> PAGE_KEYS[i % PAGE_KEYS.length] + "," + PROFILE_NAMES[i % PROFILE_NAMES.length]) |
| .collect(Collectors.toList()); |
| } |
| |
| public static List<String> getPageKeyProfileNameAddressJoin(int numMessages) { |
| return IntStream.range(0, numMessages) |
| .mapToObj(i -> PAGE_KEYS[i % PAGE_KEYS.length] + "," + PROFILE_NAMES[i % PROFILE_NAMES.length] + "," + |
| PROFILE_ZIPS[i % PROFILE_ZIPS.length] + "," + STREET_NUMS[i % STREET_NUMS.length]) |
| .collect(Collectors.toList()); |
| } |
| |
| public static List<String> getPageKeyProfileNameJoinWithNullForeignKeys(int numMessages) { |
| // All even profileId foreign keys are null |
| return IntStream.range(0, numMessages / 2) |
| .mapToObj(i -> PAGE_KEYS[(i * 2 + 1) % PAGE_KEYS.length] + "," + PROFILE_NAMES[(i * 2 + 1) % PROFILE_NAMES.length]) |
| .collect(Collectors.toList()); |
| } |
| |
| public static List<String> getPageKeyProfileNameOuterJoinWithNullForeignKeys(int numMessages) { |
| // All even profileId foreign keys are null |
| return IntStream.range(0, numMessages) |
| .mapToObj(i -> PAGE_KEYS[i % PAGE_KEYS.length] + "," + ((i % 2 == 0) ? "null" : PROFILE_NAMES[i % PROFILE_NAMES.length])) |
| .collect(Collectors.toList()); |
| } |
| |
| public static List<String> getPageKeyProfileCompanyNameJoin(int numMessages) { |
| return IntStream.range(0, numMessages) |
| .mapToObj(i -> PAGE_KEYS[i % PAGE_KEYS.length] + "," + PROFILE_NAMES[i % PROFILE_NAMES.length] + |
| "," + COMPANIES[i % COMPANIES.length]) |
| .collect(Collectors.toList()); |
| } |
| |
| public static HashMap<String, Integer> getPageKeyGroupByResult(int numMessages, Set<String> includePageKeys) { |
| HashMap<String, Integer> pageKeyCountMap = new HashMap<>(); |
| int quotient = numMessages / PAGE_KEYS.length; |
| int remainder = numMessages % PAGE_KEYS.length; |
| IntStream.range(0, PAGE_KEYS.length) |
| .map(k -> { |
| if (includePageKeys.contains(PAGE_KEYS[k])) { |
| pageKeyCountMap.put(PAGE_KEYS[k], quotient + ((k < remainder) ? 1 : 0)); |
| } |
| return k; |
| }); |
| return pageKeyCountMap; |
| } |
| |
| @Override |
| public SystemConsumer getConsumer(String systemName, Config config, MetricsRegistry registry) { |
| return new TestAvroSystemConsumer(systemName, config); |
| } |
| |
| @Override |
| public SystemProducer getProducer(String systemName, Config config, MetricsRegistry registry) { |
| return new TestAvroSystemProducer(); |
| } |
| |
| @Override |
| public SystemAdmin getAdmin(String systemName, Config config) { |
| return new SimpleSystemAdmin(config); |
| } |
| |
| private class TestAvroSystemConsumer implements SystemConsumer { |
| public static final int DEFAULT_NUM_EVENTS = 10; |
| private final int numMessages; |
| private final boolean includeNullForeignKeys; |
| private final boolean includeNullSimpleRecords; |
| private final long sleepBetweenPollsMs; |
| private final Set<SystemStreamPartition> simpleRecordSsps = new HashSet<>(); |
| private final Set<SystemStreamPartition> profileRecordSsps = new HashSet<>(); |
| private final Set<SystemStreamPartition> companyRecordSsps = new HashSet<>(); |
| private final Set<SystemStreamPartition> pageViewRecordSsps = new HashSet<>(); |
| private final Map<SystemStreamPartition, Integer> curMessagesPerSsp = new HashMap<>(); |
| |
| public TestAvroSystemConsumer(String systemName, Config config) { |
| numMessages = config.getInt(String.format("systems.%s.%s", systemName, CFG_NUM_MESSAGES), DEFAULT_NUM_EVENTS); |
| includeNullForeignKeys = config.getBoolean(String.format("systems.%s.%s", systemName, |
| CFG_INCLUDE_NULL_FOREIGN_KEYS), false); |
| includeNullSimpleRecords = config.getBoolean(String.format("systems.%s.%s", systemName, |
| CFG_INCLUDE_NULL_SIMPLE_RECORDS), false); |
| sleepBetweenPollsMs = config.getLong(String.format("systems.%s.%s", systemName, CFG_SLEEP_BETWEEN_POLLS_MS), 0); |
| } |
| |
| @Override |
| public void start() { |
| } |
| |
| @Override |
| public void stop() { |
| } |
| |
| @Override |
| public void register(SystemStreamPartition systemStreamPartition, String offset) { |
| switch (systemStreamPartition.getStream().toLowerCase()) { |
| case "simple1": |
| simpleRecordSsps.add(systemStreamPartition); |
| break; |
| case "profile": |
| profileRecordSsps.add(systemStreamPartition); |
| break; |
| case "company": |
| companyRecordSsps.add(systemStreamPartition); |
| break; |
| case "pageview": |
| pageViewRecordSsps.add(systemStreamPartition); |
| break; |
| case "complex1": |
| break; |
| case "simple2": |
| break; |
| case "simple3": |
| break; |
| default: |
| Assert.assertTrue(String.format("ssp %s is not recognized", systemStreamPartition), false); |
| break; |
| } |
| curMessagesPerSsp.put(systemStreamPartition, 0); |
| } |
| |
| @Override |
| public Map<SystemStreamPartition, List<IncomingMessageEnvelope>> poll(Set<SystemStreamPartition> set, long timeout) |
| throws InterruptedException { |
| Map<SystemStreamPartition, List<IncomingMessageEnvelope>> envelopeMap = new HashMap<>(); |
| set.forEach(ssp -> { |
| int curMessages = curMessagesPerSsp.get(ssp); |
| // We send num Messages and an end of stream message following that. |
| List<IncomingMessageEnvelope> envelopes = |
| IntStream.range(curMessages, curMessages + numMessages / 4) |
| .mapToObj(i -> i < numMessages ? new IncomingMessageEnvelope(ssp, null, getKey(i, ssp), |
| getData(i, ssp)) : IncomingMessageEnvelope.buildEndOfStreamEnvelope(ssp)) |
| .collect(Collectors.toList()); |
| envelopeMap.put(ssp, envelopes); |
| curMessagesPerSsp.put(ssp, curMessages + numMessages / 4); |
| }); |
| if (sleepBetweenPollsMs > 0) { |
| Thread.sleep(sleepBetweenPollsMs); |
| } |
| |
| return envelopeMap; |
| } |
| |
| private Object getKey(int index, SystemStreamPartition ssp) { |
| if (profileRecordSsps.contains(ssp) || companyRecordSsps.contains(ssp)) { |
| return index; // Keep this value the same as the profile/company record's "id" field. |
| } |
| return "key" + index; |
| } |
| |
| private Object getData(int index, SystemStreamPartition ssp) { |
| if (simpleRecordSsps.contains(ssp)) { |
| return createSimpleRecord(index, includeNullSimpleRecords); |
| } else if (profileRecordSsps.contains(ssp)) { |
| return createProfileRecord(index); |
| } else if (companyRecordSsps.contains(ssp)) { |
| return createCompanyRecord(index); |
| } else if (pageViewRecordSsps.contains(ssp)) { |
| return createPageViewRecord(index); |
| } else { |
| return createComplexRecord(index); |
| } |
| } |
| |
| private Object createSimpleRecord(int index, boolean includeNullRecords) { |
| if (includeNullRecords && index % NULL_RECORD_FREQUENCY == 0) { |
| return null; |
| } |
| |
| GenericRecord record = new GenericData.Record(SimpleRecord.SCHEMA$); |
| record.put("id", index); |
| record.put("name", "Name" + index); |
| return record; |
| } |
| |
| private Object createProfileRecord(int index) { |
| GenericRecord record = new GenericData.Record(Profile.SCHEMA$); |
| record.put("id", index); |
| record.put("name", PROFILE_NAMES[index % PROFILE_NAMES.length]); |
| record.put("address", createProfileAddressRecord(index)); |
| record.put("companyId", includeNullForeignKeys && (index % 2 == 0) ? null : index % COMPANIES.length); |
| record.put("phoneNumbers", createProfilePhoneNumbers(index % PHONE_NUMBERS.length)); |
| return record; |
| } |
| |
| private Object createProfileAddressRecord(int index) { |
| GenericRecord record = new GenericData.Record(AddressRecord.SCHEMA$); |
| record.put("streetnum", createProfileStreetNumRecord(index)); |
| record.put("zip", PROFILE_ZIPS[index % PROFILE_NAMES.length]); |
| return record; |
| } |
| |
| private Object createProfileStreetNumRecord(int index) { |
| GenericRecord record = new GenericData.Record(StreetNumRecord.SCHEMA$); |
| record.put("number", STREET_NUMS[index % STREET_NUMS.length]); |
| return record; |
| } |
| |
| private List<Object> createProfilePhoneNumbers(int index) { |
| List<Object> phoneNums = new ArrayList<>(); |
| phoneNums.add(createPhoneNumberRecord(index, Kind.Home)); |
| phoneNums.add(createPhoneNumberRecord(index, Kind.Work)); |
| phoneNums.add(createPhoneNumberRecord(index, Kind.Cell)); |
| return phoneNums; |
| } |
| |
| private Object createPhoneNumberRecord(int index, Kind kind) { |
| GenericRecord record = new GenericData.Record(PhoneNumber.SCHEMA$); |
| StringBuilder number = new StringBuilder(PHONE_NUMBERS[index]); |
| int lastCharIdx = number.length() - 1; |
| String suffix = ""; |
| switch (kind) { |
| case Home: |
| suffix = "1"; |
| break; |
| case Work: |
| suffix = "2"; |
| break; |
| case Cell: |
| suffix = "3"; |
| break; |
| } |
| number.replace(lastCharIdx, lastCharIdx + 1, suffix); |
| record.put("number", number); |
| record.put("kind", kind); |
| return record; |
| } |
| |
| private Object createCompanyRecord(int index) { |
| GenericRecord record = new GenericData.Record(Company.SCHEMA$); |
| record.put("id", index); |
| record.put("name", COMPANIES[index % COMPANIES.length]); |
| return record; |
| } |
| |
| private Object createPageViewRecord(int index) { |
| GenericRecord record = new GenericData.Record(PageView.SCHEMA$); |
| // All even profileId foreign keys are null |
| record.put("profileId", includeNullForeignKeys && (index % 2 == 0) ? null : index); |
| record.put("pageKey", PAGE_KEYS[index % PAGE_KEYS.length]); |
| return record; |
| } |
| |
| private Object createComplexRecord(int index) { |
| |
| GenericRecord record = new GenericData.Record(ComplexRecord.SCHEMA$); |
| record.put("id", index); |
| record.put("string_value", "Name" + index); |
| record.put("bytes_value", ByteBuffer.wrap("sample bytes".getBytes())); |
| record.put("float_value0", index + 0.123456f); |
| record.put("double_value", index + 0.0123456789); |
| MyFixed myFixedVar = new MyFixed(); |
| myFixedVar.bytes(DEFAULT_TRACKING_ID_BYTES); |
| record.put("fixed_value", myFixedVar); |
| record.put("bool_value", index % 2 == 0); |
| GenericData.Array<String> arrayValues = |
| new GenericData.Array<>(index, ComplexRecord.SCHEMA$.getField("array_values").schema().getTypes().get(1)); |
| arrayValues.addAll(IntStream.range(0, index).mapToObj(String::valueOf).collect(Collectors.toList())); |
| record.put("array_values", arrayValues); |
| // record.put("union_value", "unionStrValue"); |
| GenericRecord subRecord = new GenericData.Record(SubRecord.SCHEMA$); |
| subRecord.put("id", index); |
| subRecord.put("sub_values", arrayValues); |
| record.put("union_value", subRecord); |
| Map<String, String> mapValues = new HashMap<>(); |
| mapValues.put("key0", "value0"); |
| record.put("map_values", mapValues); |
| return record; |
| } |
| } |
| |
| private class TestAvroSystemProducer implements SystemProducer { |
| |
| @Override |
| public void start() { |
| } |
| |
| @Override |
| public void stop() { |
| } |
| |
| @Override |
| public void register(String source) { |
| } |
| |
| @Override |
| public void send(String source, OutgoingMessageEnvelope envelope) { |
| LOG.info("Adding message " + envelope); |
| messages.add(envelope); |
| } |
| |
| @Override |
| public void flush(String source) { |
| } |
| } |
| } |
