From 4469a739e2c1f8e54ab9c2d86bbc40f042077119 Mon Sep 17 00:00:00 2001
From: Niels Pardon <par@zurich.ibm.com>
Date: Thu, 26 Feb 2026 21:50:39 +0100
Subject: [PATCH] feat(core): add SubstraitBuilder.fieldReference(s) methods
 for joins

Signed-off-by: Niels Pardon <par@zurich.ibm.com>
---
 .../io/substrait/dsl/SubstraitBuilder.java    | 36 +++++++++++++++++++
 .../type/proto/JoinRoundtripTest.java         |  3 +-
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java
index e2a8a5eb0..72d3189d4 100644
--- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java
+++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java
@@ -41,6 +41,7 @@
 import io.substrait.type.TypeCreator;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Optional;
 import java.util.OptionalLong;
@@ -508,10 +509,45 @@ public List<FieldReference> fieldReferences(Rel input, int... indexes) {
         .collect(java.util.stream.Collectors.toList());
   }
 
+  /**
+   * Creates a field reference for a join operation by index across both left and right relations.
+   * The index spans across the combined schema of both join inputs, where fields from the left
+   * relation come first, followed by fields from the right relation.
+   *
+   * @param inputs the join inputs containing both left and right relations
+   * @param index the zero-based field index across the combined schema of both relations
+   * @return a {@link FieldReference} pointing to the specified field in the join context
+   */
+  public FieldReference fieldReference(JoinInput inputs, int index) {
+    final List<Rel> inputsList = new LinkedList<>();
+    inputsList.add(inputs.left);
+    inputsList.add(inputs.right);
+    return FieldReference.newInputRelReference(index, inputsList);
+  }
+
   public FieldReference fieldReference(List<Rel> inputs, int index) {
     return FieldReference.newInputRelReference(index, inputs);
   }
 
+  /**
+   * Creates multiple field references for a join operation by indexes across both left and right
+   * relations. Each index spans across the combined schema of both join inputs, where fields from
+   * the left relation come first, followed by fields from the right relation.
+   *
+   * @param inputs the join inputs containing both left and right relations
+   * @param indexes the zero-based field indexes across the combined schema of both relations
+   * @return a list of {@link FieldReference} objects pointing to the specified fields in the join
+   *     context
+   */
+  public List<FieldReference> fieldReferences(JoinInput inputs, int... indexes) {
+    final List<Rel> inputsList = new LinkedList<>();
+    inputsList.add(inputs.left);
+    inputsList.add(inputs.right);
+    return Arrays.stream(indexes)
+        .mapToObj(index -> fieldReference(inputsList, index))
+        .collect(java.util.stream.Collectors.toList());
+  }
+
   public List<FieldReference> fieldReferences(List<Rel> inputs, int... indexes) {
     return Arrays.stream(indexes)
         .mapToObj(index -> fieldReference(inputs, index))
diff --git a/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java
index fd230fb6d..92eedc019 100644
--- a/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java
+++ b/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java
@@ -48,12 +48,11 @@ void mergeJoin() {
 
   @Test
   void nestedLoopJoin() {
-    List<Rel> inputRels = Arrays.asList(leftTable, rightTable);
     Rel rel =
         NestedLoopJoin.builder()
             .from(
                 sb.nestedLoopJoin(
-                    __ ->
+                    inputRels ->
                         sb.equal(sb.fieldReference(inputRels, 0), sb.fieldReference(inputRels, 5)),
                     NestedLoopJoin.JoinType.INNER,
                     leftTable,