From d2617f32b4390e46dbc4c26638f64c8b0108b083 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Guillaume=20Mass=C3=A9=20=28=E9=A9=AC=E8=B5=9B=E5=8D=AB=29?=
 <masgui@gmail.com>
Date: Mon, 10 Jul 2023 13:53:28 -0400
Subject: [PATCH] .

---
 .../isthmus/NestedStructQueryTest.java        | 157 +++++++++++++++---
 1 file changed, 138 insertions(+), 19 deletions(-)

diff --git a/isthmus/src/test/java/io/substrait/isthmus/NestedStructQueryTest.java b/isthmus/src/test/java/io/substrait/isthmus/NestedStructQueryTest.java
index 3bc14335c..d98c23701 100644
--- a/isthmus/src/test/java/io/substrait/isthmus/NestedStructQueryTest.java
+++ b/isthmus/src/test/java/io/substrait/isthmus/NestedStructQueryTest.java
@@ -2,6 +2,8 @@
 
 import io.substrait.dsl.SubstraitBuilder;
 
+import io.substrait.plan.PlanProtoConverter;
+import io.substrait.plan.ProtoPlanConverter;
 import io.substrait.proto.Plan;
 import io.substrait.proto.*;
 
@@ -23,6 +25,7 @@
 
 import org.junit.jupiter.api.Test;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Map;
 import java.util.List;
@@ -34,6 +37,130 @@ private interface Function6<A, B, C, D, E, F, R> {
     public R apply(A a, B b, C c, D d, E e, F f);
   }
 
+  @FunctionalInterface
+  private interface Function2<A, B, R> {
+    public R apply(A a, B b);
+  }
+
+  @Test
+  public void nestedQuery3() throws SqlParseException {
+      /* struct:                   index
+     *   - (a) list:               0
+     *         struct:
+     *         - (b) i32             0
+     */
+
+    final Table table =
+        new AbstractTable() {
+          @Override
+          public RelDataType getRowType(RelDataTypeFactory factory) {
+            final Function2<String, RelDataType, RelDataType> struct =
+                    (field1, tpe1) -> factory.createStructType(Arrays.asList(Pair.of(field1, tpe1)));
+
+            final Function<RelDataType, RelDataType> list =
+                (tpe) -> factory.createArrayType(tpe, -1);
+
+            final RelDataType i32 = factory.createSqlType(SqlTypeName.INTEGER);
+            final RelDataType string = factory.createSqlType(SqlTypeName.VARCHAR);
+
+            return struct.apply("a", struct.apply("b", i32));
+          }
+        };
+    final Schema schema =
+        new AbstractSchema() {
+          @Override
+          protected Map<String, Table> getTableMap() {
+            return Map.of("my_table", table);
+          }
+        };
+
+    final SqlToSubstrait sqlToSubstrait = new SqlToSubstrait();
+
+    final SubstraitBuilder builder = new SubstraitBuilder(extensions);
+
+      String query =
+          """
+           SELECT
+             "nested"."my_table"."a"[1]."b"
+           FROM
+             "nested"."my_table";
+           """;
+
+      Plan plan = sqlToSubstrait.execute(query, "nested", schema);
+
+      try {
+          ProtoPlanConverter converter = new ProtoPlanConverter();
+          io.substrait.plan.Plan plan2 = converter.from(plan);
+          assertPlanRoundrip(plan2);
+      } catch (IOException e) {
+          throw new RuntimeException(e);
+      }
+
+  }
+
+  @Test
+  public void nestedQuery2() throws SqlParseException {
+      /* struct:                   index
+     *   - (a) list:               0
+     *         struct:
+     *         - (b) i32             0
+     */
+
+    final Table table =
+        new AbstractTable() {
+          @Override
+          public RelDataType getRowType(RelDataTypeFactory factory) {
+            final Function2<String, RelDataType, RelDataType> struct =
+                    (field1, tpe1) -> factory.createStructType(Arrays.asList(Pair.of(field1, tpe1)));
+
+            final Function<RelDataType, RelDataType> list =
+                (tpe) -> factory.createArrayType(tpe, -1);
+
+            final RelDataType i32 = factory.createSqlType(SqlTypeName.INTEGER);
+            final RelDataType string = factory.createSqlType(SqlTypeName.VARCHAR);
+
+
+            return struct.apply(
+                "a", list.apply(
+                  struct.apply(
+                    "b", i32
+                  )
+                )
+            );
+          }
+        };
+    final Schema schema =
+        new AbstractSchema() {
+          @Override
+          protected Map<String, Table> getTableMap() {
+            return Map.of("my_table", table);
+          }
+        };
+
+    final SqlToSubstrait sqlToSubstrait = new SqlToSubstrait();
+
+    final SubstraitBuilder builder = new SubstraitBuilder(extensions);
+
+      String query =
+          """
+           SELECT
+             "nested"."my_table"."a"[1]."b"
+           FROM
+             "nested"."my_table";
+           """;
+
+      Plan plan = sqlToSubstrait.execute(query, "nested", schema);
+
+      try {
+          ProtoPlanConverter converter = new ProtoPlanConverter();
+          io.substrait.plan.Plan plan2 = converter.from(plan);
+          assertPlanRoundrip(plan2);
+      } catch (IOException e) {
+          throw new RuntimeException(e);
+      }
+
+  }
+
   @Test
   public void nestedQuery() throws SqlParseException {
 
@@ -103,35 +230,27 @@ protected Map<String, Table> getTableMap() {
 
     final SubstraitBuilder builder = new SubstraitBuilder(extensions);
 
-    try {
       String query =
           """
            SELECT
-             "nested"."my_table"."a"."b",
-             "nested"."my_table"."a"."c"[1]."d",
-             "nested"."my_table"."h"['key1']['key2']."x"
+             -- "nested"."my_table"."a"."b",
+             "nested"."my_table"."a"."c"[1]."d"
+             -- "nested"."my_table"."h"['key1']['key2']."x"
            FROM
              "nested"."my_table";
            """;
 
-      Plan protoPlan1 = sqlToSubstrait.execute(query, "nested", schema);
-
-      List<Expression> expressions =
-          protoPlan1.getRelations(0).getRoot().getInput().getProject().getExpressionsList();
 
-      System.out.println(expressions);
 
-      // try {
-      //   System.out.println("yolo");
-      //   // System.out.println(JsonFormat.printer().includingDefaultValueFields().print(out));
-      // } catch (InvalidProtocolBufferException e) {
-      //   throw new RuntimeException(e);
-      // }
+      Plan plan = sqlToSubstrait.execute(query, "nested", schema);
 
-    } catch (UnsupportedOperationException op) {
-      op.printStackTrace();
-    }
+      try {
+          ProtoPlanConverter converter = new ProtoPlanConverter();
+          io.substrait.plan.Plan plan2 = converter.from(plan);
+          assertPlanRoundrip(plan2);
+      } catch (IOException e) {
+          throw new RuntimeException(e);
+      }
 
-    // System.out.println(plan);
   }
 }