@Test
public void testSplitIntoSingleRecordBundles() throws Exception {
  PipelineOptions options = PipelineOptionsFactory.create();
  SyntheticSourceOptions sourceOptions = new SyntheticSourceOptions();
  sourceOptions.numRecords = 10;
  sourceOptions.setSeed(123456);
  sourceOptions.bundleSizeDistribution = fromRealDistribution(new ConstantRealDistribution(1.0));
  sourceOptions.forceNumInitialBundles = 10;
  SyntheticBoundedSource source = new SyntheticBoundedSource(sourceOptions);
  List<SyntheticBoundedSource> sources = source.split(42L, options);
  for (SyntheticBoundedSource recordSource : sources) {
    recordSource.validate();
    assertEquals(1, recordSource.getEndOffset() - recordSource.getStartOffset());
  }
  SourceTestUtils.assertSourcesEqualReferenceSource(source, sources, options);
}
 

@Override
public Object doWork(Object first) throws IOException{
  if(null == first){
    throw new IOException(String.format(Locale.ROOT,"Invalid expression %s - null found for the first value",toExpression(constructingFactory)));
  }

  Number constant = (Number)first;

  return new ConstantRealDistribution(constant.doubleValue());
}
 

@Test
public void bundlesShouldBeEvenForConstDistribution() {
  long expectedBundleSize = 2;
  options.bundleSizeDistribution = fromRealDistribution(new ConstantRealDistribution(2));
  splitter = new BundleSplitter(options);

  List<OffsetRange> bundleSizes = splitter.getBundleSizes(4, 0, options.numRecords);

  bundleSizes.stream()
      .map(range -> range.getTo() - range.getFrom())
      .forEach(size -> assertEquals(expectedBundleSize, size.intValue()));
}
 

@Test
public void bundleSizesShouldBeProportionalToTheOneSuggestedInBundleSizeDistribution() {
  long expectedBundleSize = 4;
  options.bundleSizeDistribution = fromRealDistribution(new ConstantRealDistribution(2));
  options.numRecords = 16;
  splitter = new BundleSplitter(options);

  List<OffsetRange> bundleSizes = splitter.getBundleSizes(4, 0, options.numRecords);

  bundleSizes.stream()
      .map(range -> range.getTo() - range.getFrom())
      .forEach(size -> assertEquals(expectedBundleSize, size.intValue()));
}
 

@Test
public void consequentBundlesShouldHaveTheSameRangeEndAndStart() {
  int desiredNumberOfBundles = 2;
  options.bundleSizeDistribution = fromRealDistribution(new ConstantRealDistribution(2));
  splitter = new BundleSplitter(options);

  List<OffsetRange> bundleSizes =
      splitter.getBundleSizes(desiredNumberOfBundles, 0, options.numRecords);

  assertEquals(bundleSizes.get(0).getTo(), bundleSizes.get(1).getFrom());
  assertEquals(bundleSizes.get(0).getTo(), bundleSizes.get(1).getFrom());
  assertEquals(desiredNumberOfBundles, bundleSizes.size());
}
 

@Test
public void testSyntheticStepWithPreservingInputKeyDistribution() throws Exception {
  SyntheticStep.Options options =
      SyntheticTestUtils.optionsFromString(
          "{\"outputRecordsPerInputRecord\": 2,"
              + " \"preservesInputKeyDistribution\": true,"
              + "\"keySizeBytes\": 10,"
              + "\"valueSizeBytes\": 20,"
              + "\"numHotKeys\": 3,"
              + "\"hotKeyFraction\": 0.3,"
              + "\"seed\": 123456}",
          SyntheticStep.Options.class);
  options.delayDistribution =
      SyntheticOptions.fromRealDistribution(new ConstantRealDistribution(10));

  PCollection<byte[]> result =
      p.apply(
              Create.of(
                  ImmutableList.of(
                      KV.of(intToByteArray(1), intToByteArray(11)),
                      KV.of(intToByteArray(2), intToByteArray(22)),
                      KV.of(intToByteArray(3), intToByteArray(33)))))
          .apply(ParDo.of(new SyntheticStep(options)))
          .apply(Keys.create());

  List<byte[]> expected =
      ImmutableList.of(
          intToByteArray(1),
          intToByteArray(1),
          intToByteArray(2),
          intToByteArray(2),
          intToByteArray(3),
          intToByteArray(3));
  PAssert.that(result).containsInAnyOrder(expected);
  p.run().waitUntilFinish();
}
 

@Test
public void testSyntheticStepWithoutPreservingInputKeyDistribution() throws Exception {
  SyntheticStep.Options options =
      SyntheticTestUtils.optionsFromString(
          "{\"outputRecordsPerInputRecord\": 2,"
              + " \"preservesInputKeyDistribution\": false,"
              + "\"keySizeBytes\": 10,"
              + "\"valueSizeBytes\": 20,"
              + "\"numHotKeys\": 3,"
              + "\"hotKeyFraction\": 0.3,"
              + "\"seed\": 123456}",
          SyntheticStep.Options.class);
  options.delayDistribution =
      SyntheticOptions.fromRealDistribution(new ConstantRealDistribution(10));

  PCollection<KV<byte[], byte[]>> result =
      p.apply(Create.of(ImmutableList.of(KV.of(intToByteArray(1), intToByteArray(11)))))
          .apply(ParDo.of(new SyntheticStep(options)));

  PAssert.that(result)
      .satisfies(
          (Iterable<KV<byte[], byte[]>> input) -> {
            int count = 0;
            for (KV<byte[], byte[]> elm : input) {
              count += 1;
              assertEquals(10, elm.getKey().length);
              assertEquals(20, elm.getValue().length);
            }
            assertEquals(2, count);
            return null;
          });
  p.run().waitUntilFinish();
}
 

/**
 * The within-bin smoothing kernel. Returns a Gaussian distribution
 * parameterized by {@code bStats}, unless the bin contains only one
 * observation, in which case a constant distribution is returned.
 *
 * @param bStats summary statistics for the bin
 * @return within-bin kernel parameterized by bStats
 */
protected RealDistribution getKernel(SummaryStatistics bStats) {
    if (bStats.getN() == 1) {
        return new ConstantRealDistribution(bStats.getMean());
    } else {
        return new NormalDistribution(randomData.getRandomGenerator(),
            bStats.getMean(), bStats.getStandardDeviation(),
            NormalDistribution.DEFAULT_INVERSE_ABSOLUTE_ACCURACY);
    }
}
 

/**
 * The within-bin smoothing kernel. Returns a Gaussian distribution
 * parameterized by {@code bStats}, unless the bin contains only one
 * observation, in which case a constant distribution is returned.
 *
 * @param bStats summary statistics for the bin
 * @return within-bin kernel parameterized by bStats
 */
protected RealDistribution getKernel(SummaryStatistics bStats) {
    if (bStats.getN() == 1) {
        return new ConstantRealDistribution(bStats.getMean());
    } else {
        return new NormalDistribution(randomData.getRandomGenerator(),
            bStats.getMean(), bStats.getStandardDeviation(),
            NormalDistribution.DEFAULT_INVERSE_ABSOLUTE_ACCURACY);
    }
}
 

@Override
protected RealDistribution getKernel(SummaryStatistics bStats) {
    return new ConstantRealDistribution(bStats.getMean());
}
 

@Override
protected RealDistribution getKernel(SummaryStatistics bStats) {
    return new ConstantRealDistribution(bStats.getMean());
}