Skip to content

Commit eb31c8d

Browse files
gem-neo4jgem-neo4j
committed
[lYfZxdRz] Migrate procedures and functions from Core to Cypher 25
1 parent b24d774 commit eb31c8d

File tree

6 files changed

+1061
-0
lines changed

6 files changed

+1061
-0
lines changed

‎extended/src/main/java/apoc/export/arrow/ExportArrow.java‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
/*
2+
* Copyright (c) "Neo4j"
3+
* Neo4j Sweden AB [http://neo4j.com]
4+
*
5+
* This file is part of Neo4j.
6+
*
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
*/
19+
package apoc.export.arrow;
20+
21+
import apoc.Pools;
22+
import apoc.export.util.NodesAndRelsSubGraph;
23+
import apoc.result.ByteArrayResult;
24+
import apoc.result.ExportProgressInfo;
25+
import apoc.result.VirtualGraph;
26+
import org.neo4j.cypher.export.DatabaseSubGraph;
27+
import org.neo4j.cypher.export.SubGraph;
28+
import org.neo4j.graphdb.GraphDatabaseService;
29+
import org.neo4j.graphdb.Node;
30+
import org.neo4j.graphdb.Relationship;
31+
import org.neo4j.graphdb.Result;
32+
import org.neo4j.graphdb.Transaction;
33+
import org.neo4j.kernel.api.QueryLanguage;
34+
import org.neo4j.kernel.api.procedure.QueryLanguageScope;
35+
import org.neo4j.logging.Log;
36+
import org.neo4j.procedure.Context;
37+
import org.neo4j.procedure.Description;
38+
import org.neo4j.procedure.Name;
39+
import org.neo4j.procedure.NotThreadSafe;
40+
import org.neo4j.procedure.Procedure;
41+
import org.neo4j.procedure.TerminationGuard;
42+
43+
import java.util.Collection;
44+
import java.util.Collections;
45+
import java.util.Map;
46+
import java.util.stream.Stream;
47+
48+
public class ExportArrow {
49+
50+
@Context
51+
public Transaction tx;
52+
53+
@Context
54+
public GraphDatabaseService db;
55+
56+
@Context
57+
public Pools pools;
58+
59+
@Context
60+
public Log logger;
61+
62+
@Context
63+
public TerminationGuard terminationGuard;
64+
65+
@NotThreadSafe
66+
@Procedure(name = "apoc.export.arrow.stream.all", deprecatedBy = "This procedure is being moved to APOC Extended.")
67+
@Deprecated
68+
@QueryLanguageScope(scope = {QueryLanguage.CYPHER_25})
69+
@Description("Exports the full database as an arrow byte array.")
70+
public Stream<ByteArrayResult> all(
71+
@Name(value = "config", defaultValue = "{}", description = "{ batchSize = 2000 :: INTEGER }")
72+
Map<String, Object> config) {
73+
return new ExportArrowService(db, pools, terminationGuard, logger)
74+
.stream(new DatabaseSubGraph(tx), new ArrowConfig(config));
75+
}
76+
77+
@NotThreadSafe
78+
@Procedure(
79+
name = "apoc.export.arrow.stream.graph",
80+
deprecatedBy = "This procedure is being moved to APOC Extended.")
81+
@Deprecated
82+
@QueryLanguageScope(scope = {QueryLanguage.CYPHER_25})
83+
@Description("Exports the given graph as an arrow byte array.")
84+
public Stream<ByteArrayResult> graph(
85+
@Name(value = "graph", description = "The graph to export.") Object graph,
86+
@Name(value = "config", defaultValue = "{}", description = "{ batchSize = 2000 :: INTEGER }")
87+
Map<String, Object> config) {
88+
final SubGraph subGraph;
89+
if (graph instanceof Map) {
90+
Map<String, Object> mGraph = (Map<String, Object>) graph;
91+
if (!mGraph.containsKey("nodes")) {
92+
throw new IllegalArgumentException(
93+
"Graph Map must contains `nodes` field and `relationships` optionally");
94+
}
95+
subGraph = new NodesAndRelsSubGraph(
96+
tx, (Collection<Node>) mGraph.get("nodes"), (Collection<Relationship>) mGraph.get("relationships"));
97+
} else if (graph instanceof VirtualGraph) {
98+
VirtualGraph vGraph = (VirtualGraph) graph;
99+
subGraph = new NodesAndRelsSubGraph(tx, vGraph.nodes(), vGraph.relationships());
100+
} else {
101+
throw new IllegalArgumentException("Supported inputs are VirtualGraph, Map");
102+
}
103+
return new ExportArrowService(db, pools, terminationGuard, logger).stream(subGraph, new ArrowConfig(config));
104+
}
105+
106+
@NotThreadSafe
107+
@Procedure(
108+
name = "apoc.export.arrow.stream.query",
109+
deprecatedBy = "This procedure is being moved to APOC Extended.")
110+
@Deprecated
111+
@QueryLanguageScope(scope = {QueryLanguage.CYPHER_25})
112+
@Description("Exports the given Cypher query as an arrow byte array.")
113+
public Stream<ByteArrayResult> query(
114+
@Name(value = "query", description = "The query used to collect the data for export.") String query,
115+
@Name(value = "config", defaultValue = "{}", description = "{ batchSize = 2000 :: INTEGER }")
116+
Map<String, Object> config) {
117+
Map<String, Object> params = config == null
118+
? Collections.emptyMap()
119+
: (Map<String, Object>) config.getOrDefault("params", Collections.emptyMap());
120+
Result result = tx.execute(query, params);
121+
return new ExportArrowService(db, pools, terminationGuard, logger).stream(result, new ArrowConfig(config));
122+
}
123+
124+
@NotThreadSafe
125+
@Procedure(name = "apoc.export.arrow.all", deprecatedBy = "This procedure is being moved to APOC Extended.")
126+
@Deprecated
127+
@QueryLanguageScope(scope = {QueryLanguage.CYPHER_25})
128+
@Description("Exports the full database as an arrow file.")
129+
public Stream<ExportProgressInfo> all(
130+
@Name(value = "file", description = "The name of the file to export the data to.") String fileName,
131+
@Name(value = "config", defaultValue = "{}", description = "{ batchSize = 2000 :: INTEGER }")
132+
Map<String, Object> config) {
133+
return new ExportArrowService(db, pools, terminationGuard, logger)
134+
.file(fileName, new DatabaseSubGraph(tx), new ArrowConfig(config));
135+
}
136+
137+
@NotThreadSafe
138+
@Procedure(name = "apoc.export.arrow.graph", deprecatedBy = "This procedure is being moved to APOC Extended.")
139+
@Deprecated
140+
@QueryLanguageScope(scope = {QueryLanguage.CYPHER_25})
141+
@Description("Exports the given graph as an arrow file.")
142+
public Stream<ExportProgressInfo> graph(
143+
@Name(value = "file", description = "The name of the file to export the data to.") String fileName,
144+
@Name(value = "graph", description = "The graph to export.") Object graph,
145+
@Name(value = "config", defaultValue = "{}", description = "{ batchSize = 2000 :: INTEGER }")
146+
Map<String, Object> config) {
147+
final SubGraph subGraph;
148+
if (graph instanceof Map) {
149+
Map<String, Object> mGraph = (Map<String, Object>) graph;
150+
if (!mGraph.containsKey("nodes")) {
151+
throw new IllegalArgumentException(
152+
"Graph Map must contains `nodes` field and `relationships` optionally");
153+
}
154+
subGraph = new NodesAndRelsSubGraph(
155+
tx, (Collection<Node>) mGraph.get("nodes"), (Collection<Relationship>) mGraph.get("relationships"));
156+
} else if (graph instanceof VirtualGraph) {
157+
VirtualGraph vGraph = (VirtualGraph) graph;
158+
subGraph = new NodesAndRelsSubGraph(tx, vGraph.nodes(), vGraph.relationships());
159+
} else {
160+
throw new IllegalArgumentException("Supported inputs are VirtualGraph, Map");
161+
}
162+
return new ExportArrowService(db, pools, terminationGuard, logger)
163+
.file(fileName, subGraph, new ArrowConfig(config));
164+
}
165+
166+
@NotThreadSafe
167+
@Procedure(name = "apoc.export.arrow.query", deprecatedBy = "This procedure is being moved to APOC Extended.")
168+
@Deprecated
169+
@QueryLanguageScope(scope = {QueryLanguage.CYPHER_25})
170+
@Description("Exports the results from the given Cypher query as an arrow file.")
171+
public Stream<ExportProgressInfo> query(
172+
@Name(value = "file", description = "The name of the file to which the data will be exported.")
173+
String fileName,
174+
@Name(value = "query", description = "The query to use to collect the data for export.") String query,
175+
@Name(value = "config", defaultValue = "{}", description = "{ batchSize = 2000 :: INTEGER }")
176+
Map<String, Object> config) {
177+
Map<String, Object> params = config == null
178+
? Collections.emptyMap()
179+
: (Map<String, Object>) config.getOrDefault("params", Collections.emptyMap());
180+
Result result = tx.execute(query, params);
181+
return new ExportArrowService(db, pools, terminationGuard, logger)
182+
.file(fileName, result, new ArrowConfig(config));
183+
}
184+
}

‎extended/src/main/java/apoc/load/LoadArrow.java‎

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
/*
2+
* Copyright (c) "Neo4j"
3+
* Neo4j Sweden AB [http://neo4j.com]
4+
*
5+
* This file is part of Neo4j.
6+
*
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
*/
19+
package apoc.load;
20+
21+
import apoc.result.LoadDataMapResult;
22+
import apoc.util.FileUtils;
23+
import apoc.util.JsonUtil;
24+
import apoc.util.Util;
25+
import org.apache.arrow.memory.RootAllocator;
26+
import org.apache.arrow.vector.BitVector;
27+
import org.apache.arrow.vector.DateMilliVector;
28+
import org.apache.arrow.vector.FieldVector;
29+
import org.apache.arrow.vector.VectorSchemaRoot;
30+
import org.apache.arrow.vector.ipc.ArrowFileReader;
31+
import org.apache.arrow.vector.ipc.ArrowReader;
32+
import org.apache.arrow.vector.ipc.ArrowStreamReader;
33+
import org.apache.arrow.vector.util.Text;
34+
import org.neo4j.graphdb.security.URLAccessChecker;
35+
import org.neo4j.graphdb.security.URLAccessValidationError;
36+
import org.neo4j.kernel.api.QueryLanguage;
37+
import org.neo4j.kernel.api.procedure.QueryLanguageScope;
38+
import org.neo4j.procedure.Context;
39+
import org.neo4j.procedure.Description;
40+
import org.neo4j.procedure.Name;
41+
import org.neo4j.procedure.Procedure;
42+
import org.neo4j.values.storable.Values;
43+
44+
import java.io.ByteArrayInputStream;
45+
import java.io.IOException;
46+
import java.net.URISyntaxException;
47+
import java.nio.channels.SeekableByteChannel;
48+
import java.time.Instant;
49+
import java.time.ZoneOffset;
50+
import java.util.Collection;
51+
import java.util.HashMap;
52+
import java.util.Map;
53+
import java.util.Spliterator;
54+
import java.util.Spliterators;
55+
import java.util.concurrent.atomic.AtomicInteger;
56+
import java.util.function.Consumer;
57+
import java.util.stream.Collectors;
58+
import java.util.stream.Stream;
59+
import java.util.stream.StreamSupport;
60+
61+
public class LoadArrow {
62+
63+
@Context
64+
public URLAccessChecker urlAccessChecker;
65+
66+
private static class ArrowSpliterator extends Spliterators.AbstractSpliterator<LoadDataMapResult> {
67+
68+
private final ArrowReader reader;
69+
private final VectorSchemaRoot schemaRoot;
70+
private final AtomicInteger counter;
71+
72+
public ArrowSpliterator(ArrowReader reader, VectorSchemaRoot schemaRoot) throws IOException {
73+
super(Long.MAX_VALUE, Spliterator.ORDERED);
74+
this.reader = reader;
75+
this.schemaRoot = schemaRoot;
76+
this.counter = new AtomicInteger();
77+
this.reader.loadNextBatch();
78+
}
79+
80+
@Override
81+
public synchronized boolean tryAdvance(Consumer<? super LoadDataMapResult> action) {
82+
try {
83+
if (counter.get() >= schemaRoot.getRowCount()) {
84+
if (reader.loadNextBatch()) {
85+
counter.set(0);
86+
} else {
87+
return false;
88+
}
89+
}
90+
final Map<String, Object> row = schemaRoot.getFieldVectors().stream()
91+
.collect(
92+
HashMap::new,
93+
(map, fieldVector) -> map.put(fieldVector.getName(), read(fieldVector, counter.get())),
94+
HashMap::putAll); // please look at https://bugs.openjdk.java.net/browse/JDK-8148463
95+
counter.incrementAndGet();
96+
action.accept(new LoadDataMapResult(row));
97+
return true;
98+
} catch (Exception e) {
99+
return false;
100+
}
101+
}
102+
}
103+
104+
@Procedure(name = "apoc.load.arrow.stream", deprecatedBy = "This procedure is being moved to APOC Extended.")
105+
@Deprecated
106+
@QueryLanguageScope(scope = {QueryLanguage.CYPHER_25})
107+
@Description("Imports `NODE` and `RELATIONSHIP` values from the provided arrow byte array.")
108+
public Stream<LoadDataMapResult> stream(
109+
@Name(value = "source", description = "The data to load.") byte[] source,
110+
@Name(value = "config", defaultValue = "{}", description = "This value is never used.")
111+
Map<String, Object> config)
112+
throws IOException {
113+
RootAllocator allocator = new RootAllocator();
114+
ByteArrayInputStream inputStream = new ByteArrayInputStream(source);
115+
ArrowStreamReader streamReader = new ArrowStreamReader(inputStream, allocator);
116+
VectorSchemaRoot schemaRoot = streamReader.getVectorSchemaRoot();
117+
return StreamSupport.stream(new ArrowSpliterator(streamReader, schemaRoot), false)
118+
.onClose(() -> {
119+
Util.close(allocator);
120+
Util.close(streamReader);
121+
Util.close(schemaRoot);
122+
Util.close(inputStream);
123+
});
124+
}
125+
126+
@Procedure(name = "apoc.load.arrow", deprecatedBy = "This procedure is being moved to APOC Extended.")
127+
@Deprecated
128+
@QueryLanguageScope(scope = {QueryLanguage.CYPHER_25})
129+
@Description("Imports `NODE` and `RELATIONSHIP` values from the provided arrow file.")
130+
public Stream<LoadDataMapResult> file(
131+
@Name(value = "file", description = "The name of the file to import data from.") String fileName,
132+
@Name(value = "config", defaultValue = "{}", description = "This value is never used.")
133+
Map<String, Object> config)
134+
throws IOException, URISyntaxException, URLAccessValidationError {
135+
final SeekableByteChannel channel = FileUtils.inputStreamFor(fileName, null, null, null, urlAccessChecker)
136+
.asChannel();
137+
RootAllocator allocator = new RootAllocator();
138+
ArrowFileReader streamReader = new ArrowFileReader(channel, allocator);
139+
VectorSchemaRoot schemaRoot = streamReader.getVectorSchemaRoot();
140+
return StreamSupport.stream(new ArrowSpliterator(streamReader, schemaRoot), false)
141+
.onClose(() -> {
142+
Util.close(allocator);
143+
Util.close(streamReader);
144+
Util.close(schemaRoot);
145+
Util.close(channel);
146+
});
147+
}
148+
149+
private static Object read(FieldVector fieldVector, int index) {
150+
if (fieldVector.isNull(index)) {
151+
return null;
152+
} else if (fieldVector instanceof DateMilliVector) {
153+
DateMilliVector fe = (DateMilliVector) fieldVector;
154+
return Instant.ofEpochMilli(fe.get(index)).atOffset(ZoneOffset.UTC);
155+
} else if (fieldVector instanceof BitVector) {
156+
BitVector fe = (BitVector) fieldVector;
157+
return fe.get(index) == 1;
158+
} else {
159+
Object object = fieldVector.getObject(index);
160+
return getObject(object);
161+
}
162+
}
163+
164+
private static Object getObject(Object object) {
165+
if (object instanceof Collection) {
166+
return ((Collection<?>) object).stream().map(LoadArrow::getObject).collect(Collectors.toList());
167+
}
168+
if (object instanceof Map) {
169+
return ((Map<String, Object>) object)
170+
.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> getObject(e.getValue())));
171+
}
172+
if (object instanceof Text) {
173+
return object.toString();
174+
}
175+
try {
176+
// we test if is a valid Neo4j type
177+
return Values.of(object);
178+
} catch (Exception e) {
179+
// otherwise we try coerce it
180+
return valueToString(object);
181+
}
182+
}
183+
184+
private static String valueToString(Object value) {
185+
return JsonUtil.writeValueAsString(value);
186+
}
187+
}

0 commit comments

Comments
 (0)