Skip to content

Commit 506da86

Browse files
nmervaillienmervaillie
committed
Move batching and concurrency to writers
1 parent 229cd27 commit 506da86

6 files changed

Lines changed: 78 additions & 59 deletions

File tree

‎src/main/java/org/neo4j/dbcopy/DataTransfer.java‎

Lines changed: 4 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,10 @@
99
import reactor.core.publisher.Flux;
1010
import reactor.core.publisher.Mono;
1111

12-
import java.util.List;
13-
1412
class DataTransfer {
1513

1614
protected static final Logger LOG = LoggerFactory.getLogger(DataTransfer.class);
1715

18-
protected static final int WRITER_CONCURRENCY = 4;
19-
2016
private final DataReader dataReader;
2117
private final DataWriter dataWriter;
2218
private final CopyOptions copyOptions;
@@ -28,21 +24,13 @@ public DataTransfer(DataReader dataReader, DataWriter dataWriter, CopyOptions co
2824
}
2925

3026
Mono<Long> copyAllNodesAndRels() {
31-
var batchSize = copyOptions.batchSize();
3227
ProgressBar nodeProgressBar = new ProgressBar("Nodes", dataReader.getTotalNodeCount());
3328
ProgressBar relationshipProgressBar = new ProgressBar("Relationships", dataReader.getTotalRelationshipCount());
3429

35-
return readNodes()
36-
// ideally we should filter out properties to exclude here
37-
// but the nodes are immutable and that would require duplicating the node data structure here
38-
// which I don't want to do (yet)
39-
.buffer(batchSize)
40-
.doOnNext(batch -> nodeProgressBar.updateProgress(batch.size()))
41-
.flatMap(this::writeNodes, WRITER_CONCURRENCY)
42-
.then(readRels()
43-
.buffer(batchSize)
44-
.doOnNext(batch -> relationshipProgressBar.updateProgress(batch.size()))
45-
.flatMap(this::writeRels, 1)
30+
return dataWriter.writeNodes(readNodes(), copyOptions)
31+
.doOnNext(written -> nodeProgressBar.updateProgress(written.intValue()))
32+
.then(dataWriter.writeRelationships(readRels(), copyOptions)
33+
.doOnNext(written -> relationshipProgressBar.updateProgress(written.intValue()))
4634
.reduce(0L, Long::sum)
4735
)
4836
.doOnSuccess(it -> LOG.info("Relationships writing complete - {} relationships written", it));
@@ -55,13 +43,4 @@ private Flux<Node> readNodes() {
5543
private Flux<Relationship> readRels() {
5644
return dataReader.readRelationships();
5745
}
58-
59-
private Mono<Void> writeNodes(List<Node> nodes) {
60-
return dataWriter.writeNodes(nodes, copyOptions);
61-
}
62-
63-
private Mono<Long> writeRels(List<Relationship> relationships) {
64-
return dataWriter.writeRelationships(relationships, copyOptions);
65-
}
66-
6746
}

‎src/main/java/org/neo4j/dbcopy/bolt/BoltWriter.java‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static org.neo4j.driver.Values.parameters;
2323

2424
public class BoltWriter implements DataWriter {
25+
private static final int NODE_WRITE_CONCURRENCY = 4;
2526

2627
private final Driver targetDriver;
2728
private final String targetDatabase;
@@ -32,9 +33,23 @@ public BoltWriter(Driver targetDriver, String targetDatabase) {
3233
this.targetDatabase = targetDatabase;
3334
}
3435

35-
@SuppressWarnings("deprecation")
3636
@Override
37-
public Mono<Void> writeNodes(List<Node> nodes, CopyOptions copyOptions) {
37+
public Flux<Long> writeNodes(Flux<Node> nodes, CopyOptions copyOptions) {
38+
return nodes.buffer(copyOptions.batchSize())
39+
.flatMap(batch -> writeNodeBatch(batch, copyOptions), NODE_WRITE_CONCURRENCY);
40+
}
41+
42+
@Override
43+
public Flux<Long> writeRelationships(Flux<Relationship> relationships, CopyOptions copyOptions) {
44+
return relationships.buffer(copyOptions.batchSize())
45+
.concatMap(batch -> writeRelationshipBatch(batch, copyOptions));
46+
}
47+
48+
@SuppressWarnings("deprecation")
49+
private Mono<Long> writeNodeBatch(List<Node> nodes, CopyOptions copyOptions) {
50+
if (nodes.isEmpty()) {
51+
return Mono.empty();
52+
}
3853
return Flux.usingWhen(Mono.fromSupplier(getRxSession()),
3954
session -> session.executeWrite(tx -> {
4055
List<Map<String, Object>> nodeData = nodes.stream()
@@ -54,12 +69,14 @@ RETURN sourceNodeId, id(n) as targetNodeId""",
5469
}),
5570
ReactiveSession::close)
5671
.doOnNext(record -> nodeIdMappings.put(record.get("sourceNodeId").asLong(), record.get("targetNodeId").asLong()))
57-
.then();
72+
.count();
5873
}
5974

60-
@Override
6175
@SuppressWarnings("deprecation")
62-
public Mono<Long> writeRelationships(List<Relationship> relationships, CopyOptions copyOptions) {
76+
private Mono<Long> writeRelationshipBatch(List<Relationship> relationships, CopyOptions copyOptions) {
77+
if (relationships.isEmpty()) {
78+
return Mono.empty();
79+
}
6380
var relData = relationships.stream()
6481
.map(rel -> Map.of(
6582
"s", mappedNodeId(rel.startNodeId()),
@@ -77,7 +94,7 @@ RETURN count(*)""",
7794
parameters("inputList", relData))).flatMapMany(ReactiveResult::records)),
7895
ReactiveSession::close)
7996
.map(record -> record.get(0).asLong())
80-
.reduce(0L, Long::sum);
97+
.single(0L);
8198
}
8299

83100
private long mappedNodeId(long sourceNodeId) {

‎src/main/java/org/neo4j/dbcopy/io/DataWriter.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,20 @@
33
import org.neo4j.dbcopy.CopyOptions;
44
import org.neo4j.driver.types.Node;
55
import org.neo4j.driver.types.Relationship;
6-
import reactor.core.publisher.Mono;
6+
import reactor.core.publisher.Flux;
77

8-
import java.util.List;
98

109
public interface DataWriter extends AutoCloseable {
1110

12-
Mono<Void> writeNodes(List<Node> nodes, CopyOptions copyOptions);
11+
/**
12+
* returns a flux of updated node count
13+
*/
14+
Flux<Long> writeNodes(Flux<Node> nodes, CopyOptions copyOptions);
1315

14-
Mono<Long> writeRelationships(List<Relationship> relationships, CopyOptions copyOptions);
16+
/**
17+
* returns a flux of updated relationships count
18+
*/
19+
Flux<Long> writeRelationships(Flux<Relationship> relationships, CopyOptions copyOptions);
1520

1621
@Override
1722
default void close() throws Exception {

‎src/main/java/org/neo4j/dbcopy/json/JsonDataWriter.java‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import org.neo4j.driver.Value;
99
import org.neo4j.driver.Values;
1010
import org.neo4j.driver.types.*;
11+
import reactor.core.publisher.Flux;
1112
import reactor.core.publisher.Mono;
1213

1314
import java.io.IOException;
@@ -45,21 +46,28 @@ public JsonDataWriter(Path path) {
4546
}
4647

4748
@Override
48-
public Mono<Void> writeNodes(List<Node> nodes, CopyOptions copyOptions) {
49-
return Mono.fromRunnable(() -> {
50-
try {
51-
for (Node node : nodes) {
52-
writeNode(node);
53-
}
54-
json.flush();
55-
} catch (IOException e) {
56-
throw new UncheckedIOException("Failed to write node batch to JSON export", e);
57-
}
58-
});
49+
public Flux<Long> writeNodes(Flux<Node> nodes, CopyOptions copyOptions) {
50+
return nodes.buffer(copyOptions.batchSize())
51+
.concatMap(this::writeNodeBatch);
5952
}
6053

6154
@Override
62-
public Mono<Long> writeRelationships(List<Relationship> relationships, CopyOptions copyOptions) {
55+
public Flux<Long> writeRelationships(Flux<Relationship> relationships, CopyOptions copyOptions) {
56+
return relationships.buffer(copyOptions.batchSize())
57+
.concatMap(this::writeRelationshipBatch);
58+
}
59+
60+
private Mono<Long> writeNodeBatch(List<Node> nodes) {
61+
return Mono.fromCallable(() -> {
62+
for (Node node : nodes) {
63+
writeNode(node);
64+
}
65+
json.flush();
66+
return (long) nodes.size();
67+
}).onErrorMap(IOException.class, e -> new UncheckedIOException("Failed to write node batch to JSON export", e));
68+
}
69+
70+
private Mono<Long> writeRelationshipBatch(List<Relationship> relationships) {
6371
return Mono.fromCallable(() -> {
6472
for (Relationship relationship : relationships) {
6573
writeRelationship(relationship);
@@ -69,6 +77,7 @@ public Mono<Long> writeRelationships(List<Relationship> relationships, CopyOptio
6977
}).onErrorMap(IOException.class, e -> new UncheckedIOException("Failed to write relationship batch to JSON export", e));
7078
}
7179

80+
@SuppressWarnings("deprecation")
7281
private void writeNode(Node node) throws IOException {
7382
json.writeStartObject();
7483

@@ -98,6 +107,7 @@ private void writeProperties(MapAccessor nodeOrRelationship) throws IOException
98107
json.writeEndObject();
99108
}
100109

110+
@SuppressWarnings("deprecation")
101111
private void writeRelationship(Relationship rel) throws IOException {
102112
json.writeStartObject();
103113

‎src/test/java/org/neo4j/dbcopy/DataTransferTest.java‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -124,14 +124,21 @@ static class TestDataWriter implements DataWriter {
124124
List<Relationship> writtenRelationships = new ArrayList<>();
125125

126126
@Override
127-
public Mono<Void> writeNodes(List<Node> nodes, CopyOptions copyOptions) {
128-
return Mono.fromRunnable(() -> writtenNodes.addAll(nodes));
127+
public Flux<Long> writeNodes(Flux<Node> nodes, CopyOptions copyOptions) {
128+
return nodes.buffer(copyOptions.batchSize())
129+
.concatMap(batch -> Mono.fromCallable(() -> {
130+
writtenNodes.addAll(batch);
131+
return (long) batch.size();
132+
}));
129133
}
130134

131135
@Override
132-
public Mono<Long> writeRelationships(List<Relationship> relationships, CopyOptions copyOptions) {
133-
writtenRelationships.addAll(relationships);
134-
return Mono.just((long) relationships.size());
136+
public Flux<Long> writeRelationships(Flux<Relationship> relationships, CopyOptions copyOptions) {
137+
return relationships.buffer(copyOptions.batchSize())
138+
.concatMap(batch -> Mono.fromCallable(() -> {
139+
writtenRelationships.addAll(batch);
140+
return (long) batch.size();
141+
}));
135142
}
136143
}
137144
}

‎src/test/java/org/neo4j/dbcopy/json/JsonDataReaderWriterTest.java‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import org.neo4j.driver.Values;
77
import org.neo4j.driver.types.Node;
88
import org.neo4j.driver.types.Relationship;
9+
import reactor.core.publisher.Flux;
910

1011
import java.nio.file.Files;
1112
import java.nio.file.Path;
@@ -45,8 +46,8 @@ void should_round_trip_nodes_and_relationships() throws Exception {
4546
));
4647

4748
try (JsonDataWriter writer = new JsonDataWriter(jsonFile)) {
48-
writer.writeNodes(List.of(node), CopyOptions.DEFAULT).block();
49-
writer.writeRelationships(List.of(relationship), CopyOptions.DEFAULT).block();
49+
writer.writeNodes(Flux.just(node), CopyOptions.DEFAULT).blockLast();
50+
writer.writeRelationships(Flux.just(relationship), CopyOptions.DEFAULT).blockLast();
5051
}
5152

5253
List<Node> nodes;
@@ -81,8 +82,8 @@ void should_write_one_compact_json_record_per_line() throws Exception {
8182
Relationship relationship = new TestRelationship(7L, 1L, 1L, "KNOWS", Map.of("since", Values.value(2020L)));
8283

8384
try (JsonDataWriter writer = new JsonDataWriter(jsonFile)) {
84-
writer.writeNodes(List.of(node), CopyOptions.DEFAULT).block();
85-
writer.writeRelationships(List.of(relationship), CopyOptions.DEFAULT).block();
85+
writer.writeNodes(Flux.just(node), CopyOptions.DEFAULT).blockLast();
86+
writer.writeRelationships(Flux.just(relationship), CopyOptions.DEFAULT).blockLast();
8687
}
8788

8889
List<String> lines;
@@ -103,8 +104,8 @@ void should_round_trip_gzip_compressed_jsonl() throws Exception {
103104
Relationship relationship = new TestRelationship(7L, 1L, 1L, "KNOWS", Map.of("since", Values.value(2020L)));
104105

105106
try (JsonDataWriter writer = new JsonDataWriter(jsonFile)) {
106-
writer.writeNodes(List.of(node), CopyOptions.DEFAULT).block();
107-
writer.writeRelationships(List.of(relationship), CopyOptions.DEFAULT).block();
107+
writer.writeNodes(Flux.just(node), CopyOptions.DEFAULT).blockLast();
108+
writer.writeRelationships(Flux.just(relationship), CopyOptions.DEFAULT).blockLast();
108109
}
109110

110111
List<Node> nodes;
@@ -167,7 +168,7 @@ void should_fail_on_unsupported_export_value_type() throws Exception {
167168
));
168169

169170
try (JsonDataWriter writer = new JsonDataWriter(jsonFile)) {
170-
assertThatThrownBy(() -> writer.writeNodes(List.of(node), CopyOptions.DEFAULT).block())
171+
assertThatThrownBy(() -> writer.writeNodes(Flux.just(node), CopyOptions.DEFAULT).blockLast())
171172
.isInstanceOf(IllegalArgumentException.class)
172173
.hasMessageContaining("Unsupported property value type");
173174
}

0 commit comments

Comments
 (0)