Merge pull request #4 from codemix/value-traversals-order-by

Support .order().by() on ValueTraversals

Author: phpnode

packages/graph/src/Steps.ts

@@ -3890,7 +3890,7 @@ export interface OrderStepConfig extends StepConfig {
    * The directions to sort the results by.
    */
   directions: readonly {
-    key: string;
+    key?: string;
     direction: OrderDirection;
     nulls?: NullsOrdering;
   }[];
@@ -3903,18 +3903,14 @@ export class OrderStep extends Step<OrderStepConfig> {
 
   public *traverse(
     source: GraphSource<any>,
-    input: Iterable<TraversalPath<any, any, any>>,
+    input: Iterable<unknown>,
     _context?: QueryContext,
-  ): IterableIterator<TraversalPath<any, any, any>> {
+  ): IterableIterator<unknown> {
     const { directions } = this.config;
     const sorted = [...input].sort((a, b) => {
       for (const { key, direction, nulls } of directions) {
-        // First try to get the value as a bound variable (for aliases like UNWIND x)
-        // then fall back to property access (for expressions like ORDER BY n.name)
-        const aNode = a.get(key);
-        const bNode = b.get(key);
-        const aValue = aNode !== undefined ? aNode.value : a.property(key as never);
-        const bValue = bNode !== undefined ? bNode.value : b.property(key as never);
+        const aValue = resolveOrderValue(a, key);
+        const bValue = resolveOrderValue(b, key);
 
         // Handle null values according to nulls ordering
         const aIsNull = aValue === null || aValue === undefined;
@@ -3955,6 +3951,24 @@ export class OrderStep extends Step<OrderStepConfig> {
   }
 }
 
+function resolveOrderValue(item: unknown, key: string | undefined): unknown {
+  if (key === undefined) {
+    return item instanceof TraversalPath ? item.value : item;
+  }
+
+  if (item instanceof TraversalPath) {
+    // First try to get the value as a bound variable (for aliases like UNWIND x)
+    // then fall back to property access (for expressions like ORDER BY n.name)
+    return item.get(key)?.value ?? item.property(key as never);
+  }
+
+  if (typeof item === "object" && item !== null) {
+    return (item as Record<string, unknown>)[key];
+  }
+
+  return undefined;
+}
+
 export interface UnionStepConfig extends StepConfig {}
 
 export class UnionStep<const TSteps extends readonly Step<any>[]> extends ContainerStep<

packages/graph/src/Traversals.ts

@@ -1094,19 +1094,10 @@ export class OrderVertexTraversal<
     propertyName: TPropertyName,
     direction: OrderDirection = "asc",
   ) {
-    const steps = [...this.steps];
-    let orderStep =
-      steps.length > 0 && steps[steps.length - 1] instanceof OrderStep
-        ? steps[steps.length - 1]
-        : undefined;
-    if (!orderStep) {
-      steps.push(new OrderStep({ directions: [{ key: propertyName, direction }] }));
-    } else {
-      steps[steps.length - 1] = orderStep.clone({
-        directions: [...orderStep.config.directions, { key: propertyName, direction }],
-      });
-    }
-    return new OrderVertexTraversal<TSchema, TPath>(this.graph, steps);
+    return new OrderVertexTraversal<TSchema, TPath>(
+      this.graph,
+      appendOrderDirection(this.steps, { key: propertyName, direction }),
+    );
   }
 }
 
@@ -1130,10 +1121,80 @@ export class ValueTraversal<const TSchema extends GraphSchema, const TValue> ext
       new ValuesStep({}),
     ]);
   }
+
+  /**
+   * Order the values in the traversal.
+   */
+  public order() {
+    return new OrderValueTraversal<TSchema, TValue>(this.graph, [
+      ...this.steps,
+      new OrderStep({ directions: [] }),
+    ]);
+  }
+}
+
+export class OrderValueTraversal<
+  const TSchema extends GraphSchema,
+  const TValue,
+> extends ValueTraversal<TSchema, TValue> {
+  /**
+   * Order the values in the traversal by their natural value.
+   */
+  public by(): OrderValueTraversal<TSchema, TValue>;
+  /**
+   * Order the values in the traversal by a property on the value.
+   * @param propertyName The name of the property to order by.
+   * @param direction The direction to order by.
+   */
+  public by<const TPropertyName extends keyof GetValueTraversalProperties<TValue> & string>(
+    propertyName: TPropertyName,
+    direction?: OrderDirection,
+  ): OrderValueTraversal<TSchema, TValue>;
+  public by<const TPropertyName extends keyof GetValueTraversalProperties<TValue> & string>(
+    propertyName?: TPropertyName,
+    direction: OrderDirection = "asc",
+  ) {
+    return new OrderValueTraversal<TSchema, TValue>(
+      this.graph,
+      appendOrderDirection(this.steps, { key: propertyName, direction }),
+    );
+  }
 }
 
 type ValueOf<T> = T[keyof T];
 
+type GetValueTraversalProperties<TValue> =
+  TValue extends Element<any, any, infer TProperties, any>
+    ? TProperties
+    : TValue extends object
+      ? TValue
+      : never;
+
+function appendOrderDirection(
+  steps: readonly Step<any>[],
+  orderDirection: {
+    key?: string;
+    direction: OrderDirection;
+  },
+) {
+  const nextSteps = [...steps];
+  const orderStep =
+    nextSteps.length > 0 && nextSteps[nextSteps.length - 1] instanceof OrderStep
+      ? nextSteps[nextSteps.length - 1]
+      : undefined;
+
+  if (!orderStep) {
+    nextSteps.push(new OrderStep({ directions: [orderDirection] }));
+    return nextSteps;
+  }
+
+  nextSteps[nextSteps.length - 1] = orderStep.clone({
+    directions: [...orderStep.config.directions, orderDirection],
+  });
+
+  return nextSteps;
+}
+
 type ResolveTraversalPathProperty<
   TSchema extends GraphSchema,
   TPath,

packages/graph/src/test/ValueTraversal.test.ts

@@ -55,6 +55,56 @@ test("ValueTraversal Operations - values() extraction - values() after map extra
   expect(names.includes("Alice")).toBe(true);
 });
 
+test("ValueTraversal Operations - order() operation - order().by() sorts primitive values", () => {
+  const names = Array.from(
+    g
+      .V()
+      .hasLabel("Person")
+      .map((path) => path.value.get("name"))
+      .order()
+      .by()
+      .values(),
+  );
+
+  expect(names.length).toBeGreaterThan(0);
+  expect(names).toEqual([...names].sort());
+});
+
+test("ValueTraversal Operations - order() operation - order().by() sorts unfolded primitive values", () => {
+  const results = Array.from(
+    g
+      .V(alice.id)
+      .map(() => [3, 1, 2])
+      .unfold()
+      .order()
+      .by()
+      .values(),
+  );
+
+  expect(results).toEqual([1, 2, 3]);
+});
+
+test("ValueTraversal Operations - order() operation - order().by(property) sorts object values", () => {
+  const results = Array.from(
+    g
+      .V()
+      .hasLabel("Person")
+      .map((path) => ({
+        bucket: path.value.get("age") >= 35 ? "older" : "younger",
+        name: path.value.get("name"),
+      }))
+      .order()
+      .by("bucket")
+      .by("name")
+      .values(),
+  );
+
+  expect(results.length).toBeGreaterThan(0);
+
+  const keys = results.map(({ bucket, name }) => `${bucket}:${name}`);
+  expect(keys).toEqual([...keys].sort());
+});
+
 test("ValueTraversal Operations - values() extraction - values() on edge traversal", () => {
   const edges = Array.from(g.E().values());