(beam) branch master updated: Make Neo4JIO streamed (#29853)

[yh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

yhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 18f1b9a72ef Make Neo4JIO streamed (#29853)
18f1b9a72ef is described below

commit 18f1b9a72ef5264bbef60bd22111b62092e1eb42
Author: S. Veyrié <sy...@quansic.com>
AuthorDate: Tue Jan 2 16:35:23 2024 +0100

    Make Neo4JIO streamed (#29853)
    
    Use result.next instead of result.list that returns the full result, stored in memory
---
 .../java/org/apache/beam/sdk/io/neo4j/Neo4jIO.java | 28 ++++++++++++----------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/sdks/java/io/neo4j/src/main/java/org/apache/beam/sdk/io/neo4j/Neo4jIO.java b/sdks/java/io/neo4j/src/main/java/org/apache/beam/sdk/io/neo4j/Neo4jIO.java
index 268c6b088d8..1dd95b44c2f 100644
--- a/sdks/java/io/neo4j/src/main/java/org/apache/beam/sdk/io/neo4j/Neo4jIO.java
+++ b/sdks/java/io/neo4j/src/main/java/org/apache/beam/sdk/io/neo4j/Neo4jIO.java
@@ -814,17 +814,19 @@ public class Neo4jIO {
       // We could actually read and write here depending on the type of transaction
       // we picked.  As long as the Cypher statement returns values it's fine.
       //
-      TransactionWork<List<OutputT>> transactionWork =
+      TransactionWork<Long> transactionWork =
           transaction -> {
+            long count = 0L;
             Result result = transaction.run(cypher, parametersMap);
-            return result.list(
-                record -> {
-                  try {
-                    return rowMapper.mapRow(record);
-                  } catch (Exception e) {
-                    throw new RuntimeException("error mapping Neo4j record to row", e);
-                  }
-                });
+            while (result.hasNext()) {
+              try {
+                processContext.output(rowMapper.mapRow(result.next()));
+              } catch (Exception e) {
+                throw new RuntimeException("Error mapping Neo4J result", e);
+              }
+              count++;
+            }
+            return count;
           };
 
       if (logCypher) {
@@ -846,13 +848,13 @@ public class Neo4jIO {
       if (driverSession.session == null) {
         throw new RuntimeException("neo4j session was not initialized correctly");
       } else {
-        List<OutputT> outputs;
+        final Long count;
         if (writeTransaction) {
-          outputs = driverSession.session.writeTransaction(transactionWork, transactionConfig);
+          count = driverSession.session.writeTransaction(transactionWork, transactionConfig);
         } else {
-          outputs = driverSession.session.readTransaction(transactionWork, transactionConfig);
+          count = driverSession.session.readTransaction(transactionWork, transactionConfig);
         }
-        outputs.forEach(processContext::output);
+        LOG.debug("Retrieved " + count + " elements from Neo4J");
       }
     }
   }