Examples with Stage org.neo4j.internal.batchimport.staging.Stage used on opensource projects

Example 1 with Stage

use of org.neo4j.internal.batchimport.staging.Stage in project neo4j by neo4j.

the class ReadGroupsFromCacheStepTest method shouldProduceCompleteBatchesPerOwner.

@Test
void shouldProduceCompleteBatchesPerOwner() {
    // GIVEN
    Configuration config = Configuration.withBatchSize(Configuration.DEFAULT, 10);
    Iterator<RelationshipGroupRecord> groups = groups(new Group(1, 3), new Group(2, 3), // ^^^ perfect batch size
    new Group(3, 4), new Group(4, 2), // ^^^ slightly bigger than batch size
    new Group(5, 10), // ^^^ much bigger than batch size
    new Group(6, 35), new Group(7, 2)).iterator();
    final AtomicInteger processCounter = new AtomicInteger();
    Stage stage = new Stage(getClass().getSimpleName(), null, config, 0) {

        {
            add(new ReadGroupsFromCacheStep(control(), config, groups, 1));
            add(new VerifierStep(control(), config, processCounter));
        }
    };
    // WHEN processing the data
    ExecutionSupervisors.superviseDynamicExecution(stage);
    // THEN
    assertEquals(4, processCounter.get());
}

Example 2 with Stage

use of org.neo4j.internal.batchimport.staging.Stage in project neo4j by neo4j.

the class ReadEntityIdsStepUsingTokenIndexTest method shouldSeeRecentUpdatesRightInFrontOfExternalUpdatesPoint.

@Test
void shouldSeeRecentUpdatesRightInFrontOfExternalUpdatesPoint() throws Exception {
    // given
    long entityCount = 1_000 + random.nextInt(100);
    BitSet expectedEntityIds = new BitSet();
    BitSet seenEntityIds = new BitSet();
    try (var indexAccessor = indexAccessor()) {
        populateTokenIndex(indexAccessor, expectedEntityIds, entityCount);
        Configuration configuration = withBatchSize(DEFAULT, 100);
        Stage stage = new Stage("Test", null, configuration, 0) {

            {
                add(new ReadEntityIdsStep(control(), configuration, cursorContext -> new TokenIndexScanIdIterator(indexAccessor.newTokenReader(), new int[] { TOKEN_ID }, CursorContext.NULL), NULL, new ControlledUpdatesCheck(indexAccessor, expectedEntityIds), new AtomicBoolean(true)));
                add(new CollectEntityIdsStep(control(), configuration, seenEntityIds));
            }
        };
        // when
        stage.execute().awaitCompletion();
        // then
        assertThat(seenEntityIds).isEqualTo(expectedEntityIds);
    }
}