GML-1432: Refactor dataloaders

pyTigerGraph/gds/gds.py

@@ -938,7 +938,9 @@ def edgeNeighborLoader(
         timeout: int = 300000,
         callback_fn: Callable = None,
         reinstall_query: bool = False,
-        distributed_query: bool = False
+        distributed_query: bool = False,
+        num_machines: int = 1,
+        num_segments: int = 20
     ) -> EdgeNeighborLoader:
         """Returns an `EdgeNeighborLoader` instance.
         An `EdgeNeighborLoader` instance performs neighbor sampling from all edges in the graph in batches in the following manner:
@@ -1098,7 +1100,9 @@ def edgeNeighborLoader(
                     "delimiter": delimiter,
                     "timeout": timeout,
                     "callback_fn": callback_fn,
-                    "distributed_query": distributed_query
+                    "distributed_query": distributed_query,
+                    "num_machines": num_machines,
+                    "num_segments": num_segments
                 }
                 if self.kafkaConfig:
                     params.update(self.kafkaConfig)
@@ -1130,7 +1134,9 @@ def edgeNeighborLoader(
                 "delimiter": delimiter,
                 "timeout": timeout,
                 "callback_fn": callback_fn,
-                "distributed_query": distributed_query
+                "distributed_query": distributed_query,
+                "num_machines": num_machines,
+                "num_segments": num_segments
             }
             if self.kafkaConfig:
                 params.update(self.kafkaConfig)

pyTigerGraph/gds/gsql/dataloaders/edge_loader.gsql

@@ -1,12 +1,11 @@
 CREATE QUERY edge_loader_{QUERYSUFFIX}(
-    INT batch_size,
-    INT num_batches=1, 
-    BOOL shuffle=FALSE,
     STRING filter_by,
     SET<STRING> e_types,
     STRING delimiter,
+    BOOL shuffle=FALSE,
+    INT num_chunks=2,
     STRING kafka_address="",
-    STRING kafka_topic,
+    STRING kafka_topic="",
     INT kafka_topic_partitions=1,
     STRING kafka_max_size="104857600",
     INT kafka_timeout=300000,
@@ -41,198 +40,72 @@ CREATE QUERY edge_loader_{QUERYSUFFIX}(
       sasl_password  : SASL password for Kafka. 
       ssl_ca_location: Path to CA certificate for verifying the Kafka broker key
     */
-    TYPEDEF TUPLE<INT tmp_id, VERTEX src, VERTEX tgt> ID_Tuple;
-    INT num_vertices;
-    INT kafka_errcode;
     SumAccum<INT> @tmp_id;
-    SumAccum<STRING> @@kafka_error;
-    UINT producer;
-    MapAccum<INT, BOOL> @@edges_sampled;
-    SetAccum<VERTEX> @valid_v_out;
-    SetAccum<VERTEX> @valid_v_in;
-
-    # Initialize Kafka producer
-    IF kafka_address != "" THEN
-        producer = init_kafka_producer(
-            kafka_address, kafka_max_size, security_protocol, 
-            sasl_mechanism, sasl_username, sasl_password, ssl_ca_location,
-            ssl_certificate_location, ssl_key_location, ssl_key_password,
-            ssl_endpoint_identification_algorithm, sasl_kerberos_service_name,
-            sasl_kerberos_keytab, sasl_kerberos_principal);
-    END;
 
-    # Shuffle vertex ID if needed
     start = {ANY};
+    # Filter seeds if needed
+    seeds = SELECT s
+        FROM start:s -(e_types:e)- :t
+        WHERE filter_by is NULL OR e.getAttr(filter_by, "BOOL")
+        POST-ACCUM s.@tmp_id = getvid(s)
+        POST-ACCUM t.@tmp_id = getvid(t);
+    # Shuffle vertex ID if needed
     IF shuffle THEN
-        num_vertices = start.size();
+        INT num_vertices = seeds.size();
         res = SELECT s 
-              FROM start:s
-              POST-ACCUM s.@tmp_id = floor(rand()*num_vertices);
-    ELSE
-        res = SELECT s 
-              FROM start:s
-              POST-ACCUM s.@tmp_id = getvid(s);
-    END;
-
-    SumAccum<FLOAT> @@num_edges;
-    IF filter_by IS NOT NULL THEN
-      res = SELECT s
-            FROM start:s -(e_types:e)- :t WHERE e.getAttr(filter_by, "BOOL")
-            ACCUM 
-              IF e.isDirected() THEN # we divide by two later to correct for undirected edges being counted twice, need to count directed edges twice to get correct count
-                @@num_edges += 2
-              ELSE
-                @@num_edges += 1
-              END;
-    ELSE
-      res = SELECT s 
-            FROM start:s -(e_types:e)- :t 
-            ACCUM
-              IF e.isDirected() THEN # we divide by two later to correct for undirected edges being counted twice, need to count directed edges twice to get correct count
-                @@num_edges += 2
-              ELSE
-                @@num_edges += 1
-              END;
-    END;
-    INT batch_s;     
-    IF batch_size IS NULL THEN
-      batch_s = ceil((@@num_edges/2)/num_batches);
-    ELSE  
-      batch_s = batch_size;
+            FROM seeds:s
+            POST-ACCUM s.@tmp_id = floor(rand()*num_vertices)
+            LIMIT 1;
     END;
 
     # Generate batches
-    FOREACH batch_id IN RANGE[0, num_batches-1] DO
-        SumAccum<STRING> @@e_batch;
-        SetAccum<VERTEX> @@seeds;
-        SetAccum<VERTEX> @@targets;
-        HeapAccum<ID_Tuple> (1, tmp_id ASC) @@batch_heap;
-        @@batch_heap.resize(batch_s);
-        start = {ANY};
-        IF filter_by IS NOT NULL THEN
-            res = 
-                SELECT s
-                FROM start:s -(e_types:e)- :t
-                WHERE e.getAttr(filter_by, "BOOL") 
-                      AND 
-                        ((e.isDirected() AND ((t.@tmp_id >= s.@tmp_id AND NOT @@edges_sampled.containsKey((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id)) OR
-                                             (t.@tmp_id < s.@tmp_id AND NOT @@edges_sampled.containsKey((s.@tmp_id*s.@tmp_id)+t.@tmp_id))))
-                        OR
-                          (NOT e.isDirected() AND ((t.@tmp_id >= s.@tmp_id AND NOT @@edges_sampled.containsKey((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id)) OR
-                                                    (t.@tmp_id < s.@tmp_id AND NOT @@edges_sampled.containsKey((s.@tmp_id*s.@tmp_id)+t.@tmp_id)))
-                                              AND ((s.@tmp_id >= t.@tmp_id AND NOT @@edges_sampled.containsKey((s.@tmp_id*s.@tmp_id)+t.@tmp_id+s.@tmp_id)) OR
-                                                   (t.@tmp_id < s.@tmp_id AND NOT @@edges_sampled.containsKey((t.@tmp_id*t.@tmp_id)+s.@tmp_id)))))
-                ACCUM 
-                    IF t.@tmp_id >= s.@tmp_id THEN
-                        @@batch_heap += ID_Tuple(((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id), s, t)
-                    ELSE
-                        @@batch_heap += ID_Tuple(((s.@tmp_id*s.@tmp_id)+t.@tmp_id), s, t)
-                    END;
-
-            FOREACH elem IN @@batch_heap DO
-              SetAccum<VERTEX> @@src;
-              @@seeds += elem.src;
-              @@targets += elem.tgt;
-              @@src += elem.src;
-              src = {@@src};
-              res = SELECT s FROM src:s -(e_types:e)- :t
-                    WHERE t == elem.tgt
-                    ACCUM
-                      s.@valid_v_out += elem.tgt,
-                      t.@valid_v_in += elem.src;
-            END;
-            start = {@@seeds};
-            res =
-                SELECT s 
-                FROM start:s -(e_types:e)- :t 
-                WHERE t in @@targets AND s IN t.@valid_v_in AND t IN s.@valid_v_out
-                ACCUM 
-                    {EDGEATTRS},
-                    IF t.@tmp_id >= s.@tmp_id THEN
-                      @@edges_sampled += (((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id) -> TRUE),
-                      IF NOT e.isDirected() THEN
-                        @@edges_sampled += (((s.@tmp_id*s.@tmp_id)+t.@tmp_id+s.@tmp_id) -> TRUE)
-                      END
-                    ELSE
-                      @@edges_sampled += (((s.@tmp_id*s.@tmp_id)+t.@tmp_id) -> TRUE),
-                      IF NOT e.isDirected() THEN
-                        @@edges_sampled += (((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id) -> TRUE)
-                      END
-                    END
-                POST-ACCUM
-                    s.@valid_v_in.clear(), s.@valid_v_out.clear()
-                POST-ACCUM
-                    t.@valid_v_in.clear(), t.@valid_v_out.clear();
-        ELSE
-            res = 
-                SELECT s
-                FROM start:s -(e_types:e)- :t
-                WHERE ((e.isDirected() AND ((t.@tmp_id >= s.@tmp_id AND NOT @@edges_sampled.containsKey((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id)) OR
-                                             (t.@tmp_id < s.@tmp_id AND NOT @@edges_sampled.containsKey((s.@tmp_id*s.@tmp_id)+t.@tmp_id))))
-                          OR
-                          (NOT e.isDirected() AND ((t.@tmp_id >= s.@tmp_id AND NOT @@edges_sampled.containsKey((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id)) OR
-                                                    (t.@tmp_id < s.@tmp_id AND NOT @@edges_sampled.containsKey((s.@tmp_id*s.@tmp_id)+t.@tmp_id)))
-                                              AND ((s.@tmp_id >= t.@tmp_id AND NOT @@edges_sampled.containsKey((s.@tmp_id*s.@tmp_id)+t.@tmp_id+s.@tmp_id)) OR
-                                                   (t.@tmp_id < s.@tmp_id AND NOT @@edges_sampled.containsKey((t.@tmp_id*t.@tmp_id)+s.@tmp_id)))))
-                ACCUM 
-                    IF t.@tmp_id >= s.@tmp_id THEN
-                        @@batch_heap += ID_Tuple(((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id), s, t)
-                    ELSE
-                        @@batch_heap += ID_Tuple(((s.@tmp_id*s.@tmp_id)+t.@tmp_id), s, t)
-                    END;
-
-            FOREACH elem IN @@batch_heap DO
-              SetAccum<VERTEX> @@src;
-              @@seeds += elem.src;
-              @@targets += elem.tgt;
-              @@src += elem.src;
-              src = {@@src};
-              res = SELECT s FROM src:s -(e_types:e)- :t
-                    WHERE t == elem.tgt
-                    ACCUM
-                      s.@valid_v_out += elem.tgt,
-                      t.@valid_v_in += elem.src;
-            END;
-            start = {@@seeds};
-            res =
-                SELECT s 
-                FROM start:s -(e_types:e)- :t 
-                WHERE t in @@targets AND s IN t.@valid_v_in AND t IN s.@valid_v_out
-                ACCUM 
-                    {EDGEATTRS},
-                    IF t.@tmp_id >= s.@tmp_id THEN
-                      @@edges_sampled += (((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id) -> TRUE),
-                      IF NOT e.isDirected() THEN
-                        @@edges_sampled += (((s.@tmp_id*s.@tmp_id)+t.@tmp_id+s.@tmp_id) -> TRUE)
-                      END
-                    ELSE
-                      @@edges_sampled += (((s.@tmp_id*s.@tmp_id)+t.@tmp_id) -> TRUE),
-                      IF NOT e.isDirected() THEN
-                        @@edges_sampled += (((t.@tmp_id*t.@tmp_id)+s.@tmp_id+t.@tmp_id) -> TRUE)
-                      END
-                    END
-                POST-ACCUM
-                    s.@valid_v_in.clear(), s.@valid_v_out.clear()
-                POST-ACCUM
-                    t.@valid_v_in.clear(), t.@valid_v_out.clear();
+    # If using kafka to export 
+    IF kafka_address != "" THEN
+        SumAccum<STRING> @@kafka_error;
+
+        # Initialize Kafka producer
+        UINT producer = init_kafka_producer(
+            kafka_address, kafka_max_size, security_protocol, 
+            sasl_mechanism, sasl_username, sasl_password, ssl_ca_location,
+            ssl_certificate_location, ssl_key_location, ssl_key_password,
+            ssl_endpoint_identification_algorithm, sasl_kerberos_service_name,
+            sasl_kerberos_keytab, sasl_kerberos_principal);
+
+        FOREACH chunk IN RANGE[0, num_chunks-1] DO
+            res = SELECT s 
+                FROM seeds:s -(e_types:e)- :t
+                WHERE (filter_by is NULL OR e.getAttr(filter_by, "BOOL")) and ((s.@tmp_id + t.@tmp_id) % num_chunks == chunk)
+                ACCUM
+                    {EDGEATTRSKAFKA}
+                LIMIT 1;
         END;
-        # Export batch
-        IF kafka_address != "" THEN
-            # Write to kafka
-            kafka_errcode = write_to_kafka(producer, kafka_topic, batch_id%kafka_topic_partitions, "edge_batch_" + stringify(batch_id), @@e_batch);
-            IF kafka_errcode != 0 THEN 
-                @@kafka_error += ("Error sending edge batch " + stringify(batch_id) + ": "+ stringify(kafka_errcode) + "\n");
+
+        FOREACH i IN RANGE[0, kafka_topic_partitions-1] DO
+            INT kafka_errcode = write_to_kafka(producer, kafka_topic, i, "STOP", "");
+            IF kafka_errcode!=0 THEN 
+                @@kafka_error += ("Error sending STOP signal to topic partition " + stringify(i) + ": " + stringify(kafka_errcode) + "\n");
             END;
-        ELSE
-            # Add to response
-            PRINT @@e_batch AS edge_batch;
         END;
-    END; 
-    IF kafka_address != "" THEN
-        kafka_errcode = close_kafka_producer(producer, kafka_timeout);
-        IF kafka_errcode != 0 THEN 
+
+        INT kafka_errcode = close_kafka_producer(producer, kafka_timeout);
+        IF kafka_errcode!=0 THEN 
             @@kafka_error += ("Error shutting down Kafka producer: " + stringify(kafka_errcode) + "\n");
         END;
         PRINT @@kafka_error as kafkaError;
+    # Else return as http response
+    ELSE
+        FOREACH chunk IN RANGE[0, num_chunks-1] DO
+            ListAccum<STRING> @@e_batch;
+            res = SELECT s 
+                FROM seeds:s -(e_types:e)- :t
+                WHERE (filter_by is NULL OR e.getAttr(filter_by, "BOOL")) and ((s.@tmp_id + t.@tmp_id) % num_chunks == chunk)
+                ACCUM 
+                    {EDGEATTRSHTTP}
+                LIMIT 1;
+
+            FOREACH i IN @@e_batch DO
+                PRINT i as data_batch;
+            END;
+        END;
     END;
 }