@Test public void testGatherBundleAndSortFn() throws Exception { GatherBundleAndSortFn testFn = new GatherBundleAndSortFn(10000000, 10, 100, null); testFn.startBundle(); for (MutationGroup m : mutationGroups) { when(mockProcessContext.element()).thenReturn(m); testFn.processElement(mockProcessContext); testFn.finishBundle(mockFinishBundleContext);
@StartBundle public synchronized void startBundle() throws Exception { if (mutationsToSort == null) { initSorter(); } else { throw new IllegalStateException("Sorter should be null here"); } }
"Gather And Sort", ParDo.of( new GatherBundleAndSortFn( spec.getBatchSizeBytes(), spec.getMaxNumMutations(),
GatherBundleAndSortFn testFn = new GatherBundleAndSortFn(10000000, CELLS_PER_KEY, 3, null); testFn.startBundle(); for (MutationGroup m : mutationGroups) { when(mockProcessContext.element()).thenReturn(m); testFn.processElement(mockProcessContext); testFn.finishBundle(mockFinishBundleContext);
@ProcessElement public void processElement(ProcessContext c) throws Exception { SpannerSchema spannerSchema = c.sideInput(schemaView); MutationKeyEncoder encoder = new MutationKeyEncoder(spannerSchema); MutationGroup mg = c.element(); long groupSize = MutationSizeEstimator.sizeOf(mg); long groupCells = MutationCellCounter.countOf(spannerSchema, mg); synchronized (this) { if (((batchCells + groupCells) > maxNumMutations) || (batchSizeBytes + groupSize) > maxBatchSizeBytes) { c.output(sortAndGetList()); initSorter(); } mutationsToSort.add(KV.of(encoder.encodeTableNameAndKey(mg.primary()), encode(mg))); batchSizeBytes += groupSize; batchCells += groupCells; } } }
@FinishBundle public synchronized void finishBundle(FinishBundleContext c) throws Exception { c.output(sortAndGetList(), Instant.now(), GlobalWindow.INSTANCE); }