Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[3X06NMak] Update Scan to PartitionedScan as Scan is deprecated #498

Merged
merged 1 commit into from
Sep 29, 2023
Merged

[3X06NMak] Update Scan to PartitionedScan as Scan is deprecated #498

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions common/src/main/java/apoc/util/kernel/MultiThreadedGlobalGraphOperations.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,9 @@

import org.neo4j.internal.kernel.api.NodeCursor;
import org.neo4j.internal.kernel.api.Read;
import org.neo4j.internal.kernel.api.Scan;
import org.neo4j.internal.kernel.api.security.AccessMode;
import org.neo4j.internal.kernel.api.PartitionedScan;
import org.neo4j.internal.kernel.api.security.LoginContext;
import org.neo4j.kernel.api.ExecutionContext;
import org.neo4j.kernel.api.KernelTransaction;
import org.neo4j.kernel.impl.coreapi.InternalTransaction;
import org.neo4j.kernel.internal.GraphDatabaseAPI;
Expand All @@ -41,8 +41,8 @@ public static BatchJobResult forAllNodes(GraphDatabaseAPI db, ExecutorService ex
AtomicInteger processing = new AtomicInteger();
try ( InternalTransaction tx = db.beginTransaction( KernelTransaction.Type.EXPLICIT, LoginContext.AUTH_DISABLED ) ) {
KernelTransaction ktx = tx.kernelTransaction();
Function<Read, Scan<NodeCursor>> scanFunction = Read::allNodesScan;
Scan<NodeCursor> scan = scanFunction.apply( ktx.dataRead() );
Read dataRead = ktx.dataRead();
PartitionedScan<NodeCursor> scan = dataRead.allNodesScan(1, ktx.cursorContext());
Function<KernelTransaction,NodeCursor> cursorAllocator = ktx2 -> ktx2.cursors().allocateNodeCursor( ktx2.cursorContext() );
executorService.submit( new BatchJob( scan, batchSize, db, consumer, result, cursorAllocator, executorService, processing ) );
}
Expand Down Expand Up @@ -80,7 +80,7 @@ public long getFailures() {
}

private static class BatchJob implements Callable<Void> {
private final Scan<NodeCursor> scan;
private final PartitionedScan<NodeCursor> scan;
private final int batchSize;
private final GraphDatabaseAPI db;
private final Consumer<NodeCursor> consumer;
Expand All @@ -89,7 +89,7 @@ private static class BatchJob implements Callable<Void> {
private final ExecutorService executorService;
private final AtomicInteger processing;

public BatchJob(Scan<NodeCursor> scan, int batchSize, GraphDatabaseAPI db, Consumer<NodeCursor> consumer,
public BatchJob(PartitionedScan<NodeCursor> scan, int batchSize, GraphDatabaseAPI db, Consumer<NodeCursor> consumer,
BatchJobResult result, Function<KernelTransaction,NodeCursor> cursorAllocator, ExecutorService executorService, AtomicInteger processing ) {
this.scan = scan;
this.batchSize = batchSize;
Expand All @@ -106,17 +106,21 @@ public BatchJob(Scan<NodeCursor> scan, int batchSize, GraphDatabaseAPI db, Consu
public Void call() {
try (InternalTransaction tx = db.beginTransaction(KernelTransaction.Type.EXPLICIT, LoginContext.AUTH_DISABLED)) {
KernelTransaction ktx = tx.kernelTransaction();
ktx.acquireStatement();
ExecutionContext executionContext = ktx.createExecutionContext();
try ( NodeCursor cursor = cursorAllocator.apply( ktx )) {
if (scan.reserveBatch( cursor, batchSize, ktx.cursorContext(), AccessMode.Static.FULL )) {
while (scan.reservePartition(cursor, executionContext)) {
// Branch out so that all available threads will get saturated
executorService.submit( new BatchJob( scan, batchSize, db, consumer, result, cursorAllocator, executorService, processing ) );
executorService.submit( new BatchJob( scan, batchSize, db, consumer, result, cursorAllocator, executorService, processing ) );
executorService.submit(new BatchJob(scan, batchSize, db, consumer, result, cursorAllocator, executorService, processing));
executorService.submit(new BatchJob(scan, batchSize, db, consumer, result, cursorAllocator, executorService, processing));
while (processAndReport(cursor)) {
// just continue processing...
}
}
}
tx.commit();
executionContext.complete();
executionContext.close();
return null;
} finally {
result.batches.incrementAndGet();
Expand Down