Sample each edge type with its own limit, 10 types per request

Author: kmcginnes

packages/graph-explorer/src/connector/gremlin/fetchEdgeConnections/edgeConnectionsTemplate.test.ts

@@ -4,23 +4,26 @@ import { normalizeWithNoSpace as normalize } from "@/utils/testing";
 import edgeConnectionsTemplate from "./edgeConnectionsTemplate";
 
 describe("Gremlin > edgeConnectionsTemplate", () => {
-  it("should scan a batch of edge types and group the endpoint labels by edge type", () => {
+  it("should give each sampled edge type its own limited branch", () => {
     const template = edgeConnectionsTemplate({
       types: [createEdgeType("route"), createEdgeType("contains")],
     });
 
+    // A single limit after hasLabel() would be shared, and a dominant type would
+    // fill it before the rest were read at all.
     expect(normalize(template)).toBe(
       normalize(`
-        g.E().hasLabel('route', 'contains')
+        g.V().limit(1).union(
+          V().outE('route').limit(10000),
+          V().outE('contains').limit(10000)
+        )
           .group()
             .by(label())
             .by(
-              limit(10000)
-                .project('sourceType', 'targetType')
-                .by(outV().label())
-                .by(inV().label())
-                .dedup()
-                .fold()
+              project('s', 't')
+                .by(outV().label().fold())
+                .by(inV().label().fold())
+                .groupCount()
             )
       `),
     );
@@ -31,6 +34,6 @@ describe("Gremlin > edgeConnectionsTemplate", () => {
       types: [createEdgeType("edge'with'quotes")],
     });
 
-    expect(template).toContain("hasLabel('edge\\'with\\'quotes')");
+    expect(template).toContain("V().outE('edge\\'with\\'quotes')");
   });
 });

packages/graph-explorer/src/connector/gremlin/fetchEdgeConnections/edgeConnectionsTemplate.ts

@@ -5,40 +5,57 @@ import { DEFAULT_SAMPLE_SIZE, query } from "@/utils";
 import { fragment } from "../fragments";
 
 /**
- * Returns a Gremlin query that discovers distinct edge connection patterns for a
- * batch of edge types in a single request.
+ * Keys of the projected endpoint labels. The template writes them and the
+ * response parser reads them, so both derive from this object.
+ */
+export const projectionKeys = {
+  sourceType: "s",
+  targetType: "t",
+} as const;
+
+/**
+ * Returns a Gremlin query that counts the endpoint label combinations of up to
+ * `DEFAULT_SAMPLE_SIZE` edges of each given type, grouped by edge type.
+ *
+ * One limit after `hasLabel(A, B, ...)` would be shared, so a dominant type fills
+ * it and the rest come back empty. Each type gets its own `union()` branch
+ * instead, and on Neptune each branch is an index lookup by edge label.
+ * Mid-traversal `V()` rather than `E()`, which needs TinkerPop 3.7, and anchored
+ * on `V().limit(1)` because anchoring on `inject()` is not native on Neptune.
+ *
+ * Grouped by edge type before counting because Neptune's DFE engine cannot count
+ * one `project()` key across several full branches: two took 54s and five timed
+ * out, where grouping first handled ten in 9s.
  *
- * `g.E().hasLabel(...)` is a native, index-backed edge scan on Neptune; the
- * `group().by(label())` reduction buckets the edges by edge type, and each
- * bucket samples up to `DEFAULT_SAMPLE_SIZE` edges and projects the distinct
- * (source label, target label) pairs. The caller regroups by the returned edge
- * label.
+ * Callers send 10 types per request, so one request reads at most 100,000
+ * edges. On a db.t3.medium, 100 types in one request took 116s and left the
+ * instance refusing even a single-type sample for two minutes afterwards.
  *
- * The `limit` sits inside the `group()` value traversal, which TinkerPop runs
- * per group, so each edge type is sampled independently — no shared cap that
- * starves rarer types. It bounds the per-type label-resolution and dedup work,
- * not the initial edge scan: `group()` still enumerates every edge of the
- * batched types to bucket them. A per-type scan cap is not expressible in one
- * native TinkerPop 3.6.2 request.
+ * The endpoint labels are folded because engines disagree on what `label()`
+ * emits for a multi-label vertex. Neptune 1.4 emits one `::` composite, but
+ * 1.3.5 emits each label separately, and a bare `by(outV().label())` keeps only
+ * the first, silently dropping the vertex's other types.
  */
 export default function edgeConnectionsTemplate({
   types,
 }: {
   types: EdgeType[];
 }) {
-  const labels = types.map(fragment.identifier);
+  const limit = fragment.number(DEFAULT_SAMPLE_SIZE);
+  const branches = types.map(
+    type => `V().outE(${fragment.identifier(type)}).limit(${limit})`,
+  );
+  const keys = Object.values(projectionKeys).map(fragment.identifier);
 
   return query`
-    g.E().hasLabel(${labels.join(", ")})
+    g.V().limit(1).union(${branches.join(", ")})
       .group()
         .by(label())
         .by(
-          limit(${DEFAULT_SAMPLE_SIZE})
-            .project('sourceType', 'targetType')
-            .by(outV().label())
-            .by(inV().label())
-            .dedup()
-            .fold()
+          project(${keys.join(", ")})
+            .by(outV().label().fold())
+            .by(inV().label().fold())
+            .groupCount()
         )
   `;
 }