Fixes #2974 Arrays should be indexable from the end too

confluentinc · Jun 21, 2019 · 3872065 · 3872065
1 parent e0d1676
commit 3872065
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 9 deletions.
diff --git a/ksql-engine/src/main/java/io/confluent/ksql/codegen/SqlToJavaVisitor.java b/ksql-engine/src/main/java/io/confluent/ksql/codegen/SqlToJavaVisitor.java
@@ -633,15 +633,30 @@ protected Pair<String, Schema> visitSubscriptExpression(
           SchemaUtil.getJavaType(internalSchema).getCanonicalName();
       switch (internalSchema.type()) {
         case ARRAY:
-          return new Pair<>(
-              String.format("((%s) ((%s)%s).get((int)(%s)))",
-                  SchemaUtil.getJavaType(internalSchema.valueSchema()).getSimpleName(),
-                  internalSchemaJavaType,
-                  process(node.getBase(), context).getLeft(),
-                  process(node.getIndex(), context).getLeft()
-              ),
-              internalSchema.valueSchema()
-          );
+          final String listName = process(node.getBase(), context).getLeft();
+          if (node.getIndex().toString().startsWith("-")) {
+            return new Pair<>(
+                    String.format("((%s) ((%s)%s).get((int)((%s)%s).size()%s))",
+                            SchemaUtil.getJavaType(internalSchema.valueSchema()).getSimpleName(),
+                            internalSchemaJavaType,
+                            listName,
+                            internalSchemaJavaType,
+                            listName,
+                            process(node.getIndex(), context).getLeft()
+                    ),
+                    internalSchema.valueSchema()
+            );
+          } else {
+            return new Pair<>(
+                    String.format("((%s) ((%s)%s).get((int)(%s)))",
+                            SchemaUtil.getJavaType(internalSchema.valueSchema()).getSimpleName(),
+                            internalSchemaJavaType,
+                            listName,
+                            process(node.getIndex(), context).getLeft()
+                    ),
+                    internalSchema.valueSchema()
+            );
+          }
         case MAP:
           return new Pair<>(
               String.format("((%s) ((%s)%s).get(%s))",

diff --git a/ksql-engine/src/test/java/io/confluent/ksql/codegen/SqlToJavaVisitorTest.java b/ksql-engine/src/test/java/io/confluent/ksql/codegen/SqlToJavaVisitorTest.java
@@ -107,6 +107,18 @@ public void shouldProcessArrayExpressionCorrectly() {
         equalTo("((Double) ((java.util.List)TEST1_COL4).get((int)(Integer.parseInt(\"0\"))))"));
   }
 
+  @Test
+  public void shouldProcessArrayNegativeIndexExpressionCorrectly() {
+    final String simpleQuery = "SELECT col4[-1] FROM test1 WHERE col0 > 100;";
+    final Analysis analysis = analyzeQuery(simpleQuery, metaStore);
+
+    final String javaExpression = sqlToJavaVisitor
+            .process(analysis.getSelectExpressions().get(0));
+
+    assertThat(javaExpression,
+            equalTo("((Double) ((java.util.List)TEST1_COL4).get((int)((java.util.List)TEST1_COL4).size()-Integer.parseInt(\"1\")))"));
+  }
+
   @Test
   public void shouldProcessMapExpressionCorrectly() {
     final String simpleQuery = "SELECT col5['key1'] FROM test1 WHERE col0 > 100;";