From 17f131bdc551842438ad4ccbcfbe2e67f1d67335 Mon Sep 17 00:00:00 2001
From: Niels Pardon <par@zurich.ibm.com>
Date: Thu, 6 Mar 2025 14:12:20 +0100
Subject: [PATCH] feat(isthmus): support converting Substrait Plan.Root to
 Calcite RelRoot

Signed-off-by: Niels Pardon <par@zurich.ibm.com>
---
 .../substrait/isthmus/SubstraitToCalcite.java | 34 +++++++++++
 .../isthmus/SubstraitToCalciteTest.java       | 59 +++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 isthmus/src/test/java/io/substrait/isthmus/SubstraitToCalciteTest.java

diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java
index 729d11bee..284ef64c4 100644
--- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java
+++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitToCalcite.java
@@ -1,10 +1,12 @@
 package io.substrait.isthmus;
 
 import io.substrait.extension.SimpleExtension;
+import io.substrait.plan.Plan;
 import io.substrait.relation.AbstractRelVisitor;
 import io.substrait.relation.NamedScan;
 import io.substrait.relation.Rel;
 import io.substrait.type.NamedStruct;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -12,8 +14,14 @@
 import org.apache.calcite.jdbc.CalciteSchema;
 import org.apache.calcite.jdbc.LookupCalciteSchema;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelRoot;
+import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.calcite.rel.type.RelDataTypeFieldImpl;
+import org.apache.calcite.rel.type.RelRecordType;
 import org.apache.calcite.schema.Table;
+import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.tools.Frameworks;
 import org.apache.calcite.tools.RelBuilder;
 
@@ -99,6 +107,32 @@ public RelNode convert(Rel rel) {
     return rel.accept(converter);
   }
 
+  /**
+   * Converts a Substrait {@link Plan.Root} to a Calcite {@link RelRoot}
+   *
+   * <p>Generates a {@link RelDataType} row type with the final field names of the {@link Plan.Root}
+   * and creates a Calcite {@link RelRoot} with it.
+   *
+   * <p>TODO: revisit this code when support for WriteRel is added to substrait-java
+   *
+   * @param root {@link Plan.Root} to convert
+   * @return {@link RelRoot}
+   */
+  public RelRoot convert(Plan.Root root) {
+    RelNode input = convert(root.getInput());
+    RelDataType inputRowType = input.getRowType();
+    List<RelDataTypeField> fields = new ArrayList<>(inputRowType.getFieldCount());
+    for (RelDataTypeField field : inputRowType.getFieldList()) {
+      RelDataTypeFieldImpl renamedField =
+          new RelDataTypeFieldImpl(
+              root.getNames().get(field.getIndex()), field.getIndex(), field.getType());
+      fields.add(renamedField);
+    }
+
+    RelRoot calciteRoot = RelRoot.of(input, new RelRecordType(fields), SqlKind.SELECT);
+    return calciteRoot;
+  }
+
   private static class NamedStructGatherer extends AbstractRelVisitor<Void, RuntimeException> {
     Map<List<String>, NamedStruct> tableMap;
 
diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitToCalciteTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitToCalciteTest.java
new file mode 100644
index 000000000..111af6ba8
--- /dev/null
+++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitToCalciteTest.java
@@ -0,0 +1,59 @@
+package io.substrait.isthmus;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.substrait.plan.Plan;
+import io.substrait.plan.ProtoPlanConverter;
+import java.util.Map.Entry;
+import org.apache.calcite.adapter.tpcds.TpcdsSchema;
+import org.apache.calcite.rel.RelRoot;
+import org.junit.jupiter.api.Test;
+
+public class SubstraitToCalciteTest extends PlanTestBase {
+  @Test
+  void testConvertRoot() throws Exception {
+    SqlToSubstrait s = new SqlToSubstrait();
+    TpcdsSchema schema = new TpcdsSchema(1.0);
+
+    // single column
+    String newColName = "store_name";
+    String sql = "select s_store_name as " + newColName + " from tpcds.store";
+
+    io.substrait.proto.Plan protoPlan = s.execute(sql, "tpcds", schema);
+
+    ProtoPlanConverter protoPlanConverter = new ProtoPlanConverter(extensions);
+    Plan plan = protoPlanConverter.from(protoPlan);
+    Plan.Root root = plan.getRoots().get(0);
+
+    assertEquals(1, root.getNames().size());
+    assertEquals(newColName.toUpperCase(), root.getNames().get(0));
+
+    SubstraitToCalcite converter = new SubstraitToCalcite(extensions, typeFactory);
+    RelRoot relRoot = converter.convert(root);
+
+    assertEquals(root.getNames().size(), relRoot.fields.size());
+    for (Entry<Integer, String> field : relRoot.fields) {
+      assertEquals(root.getNames().get(field.getKey()), field.getValue());
+    }
+
+    // multiple columns
+    String storeIdColumnName = "s_store_id";
+    sql = "select " + storeIdColumnName + ", s_store_name as " + newColName + " from tpcds.store";
+    protoPlan = s.execute(sql, "tpcds", schema);
+
+    protoPlanConverter = new ProtoPlanConverter(extensions);
+    plan = protoPlanConverter.from(protoPlan);
+    root = plan.getRoots().get(0);
+
+    assertEquals(2, root.getNames().size());
+    assertEquals(storeIdColumnName.toUpperCase(), root.getNames().get(0));
+    assertEquals(newColName.toUpperCase(), root.getNames().get(1));
+
+    relRoot = converter.convert(root);
+
+    assertEquals(root.getNames().size(), relRoot.fields.size());
+    for (Entry<Integer, String> field : relRoot.fields) {
+      assertEquals(root.getNames().get(field.getKey()), field.getValue());
+    }
+  }
+}