feat: native map/array constructors (#4232)

confluentinc · Jan 17, 2020 · 3ecfaad · 3ecfaad
1 parent 7a8ba8c
commit 3ecfaad
Show file tree

Hide file tree

Showing 28 changed files with 903 additions and 166 deletions.
diff --git a/ksql-engine/src/main/java/io/confluent/ksql/engine/rewrite/ExpressionTreeRewriter.java b/ksql-engine/src/main/java/io/confluent/ksql/engine/rewrite/ExpressionTreeRewriter.java
@@ -17,13 +17,16 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableList.Builder;
+import com.google.common.collect.ImmutableMap;
 import io.confluent.ksql.execution.expression.tree.ArithmeticBinaryExpression;
 import io.confluent.ksql.execution.expression.tree.ArithmeticUnaryExpression;
 import io.confluent.ksql.execution.expression.tree.BetweenPredicate;
 import io.confluent.ksql.execution.expression.tree.BooleanLiteral;
 import io.confluent.ksql.execution.expression.tree.Cast;
 import io.confluent.ksql.execution.expression.tree.ColumnReferenceExp;
 import io.confluent.ksql.execution.expression.tree.ComparisonExpression;
+import io.confluent.ksql.execution.expression.tree.CreateArrayExpression;
+import io.confluent.ksql.execution.expression.tree.CreateMapExpression;
 import io.confluent.ksql.execution.expression.tree.CreateStructExpression;
 import io.confluent.ksql.execution.expression.tree.CreateStructExpression.Field;
 import io.confluent.ksql.execution.expression.tree.DecimalLiteral;
@@ -51,6 +54,7 @@
 import io.confluent.ksql.execution.expression.tree.Type;
 import io.confluent.ksql.execution.expression.tree.WhenClause;
 import java.util.List;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.function.BiFunction;
@@ -188,6 +192,27 @@ public Expression visitSubscriptExpression(
       return new SubscriptExpression(node.getLocation(), base, index);
     }
 
+    @Override
+    public Expression visitCreateArrayExpression(final CreateArrayExpression exp, final C context) {
+      final Builder<Expression> values = ImmutableList.builder();
+      for (Expression value : exp.getValues()) {
+        values.add(rewriter.apply(value, context));
+      }
+      return new CreateArrayExpression(exp.getLocation(), values.build());
+    }
+
+    @Override
+    public Expression visitCreateMapExpression(final CreateMapExpression exp, final C context) {
+      final ImmutableMap.Builder<Expression, Expression> map = ImmutableMap.builder();
+      for (Entry<Expression, Expression> entry : exp.getMap().entrySet()) {
+        map.put(
+            rewriter.apply(entry.getKey(), context),
+            rewriter.apply(entry.getValue(), context)
+        );
+      }
+      return new CreateMapExpression(exp.getLocation(), map.build());
+    }
+
     @Override
     public Expression visitStructExpression(final CreateStructExpression node, final C context) {
       final Builder<Field> fields = ImmutableList.builder();

diff --git a/ksql-engine/src/main/java/io/confluent/ksql/function/udf/map/AsMap.java b/ksql-engine/src/main/java/io/confluent/ksql/function/udf/map/AsMap.java
@@ -38,5 +38,4 @@ public final <T> Map<String, T> asMap(
     }
     return map;
   }
-
-}
+}
diff --git a/ksql-engine/src/test/java/io/confluent/ksql/codegen/CodeGenRunnerTest.java b/ksql-engine/src/test/java/io/confluent/ksql/codegen/CodeGenRunnerTest.java
@@ -671,6 +671,42 @@ public void shouldHandleMaps() {
         assertThat(result, is("value1"));
     }
 
+    @Test
+    public void shouldHandleCreateArray() {
+        // Given:
+        final Expression expression = analyzeQuery(
+            "SELECT ARRAY['foo', COL" + STRING_INDEX1 + "] FROM codegen_test EMIT CHANGES;", metaStore)
+            .getSelectExpressions()
+            .get(0)
+            .getExpression();
+
+        // When:
+        final Object result = codeGenRunner
+            .buildCodeGenFromParseTree(expression, "Array")
+            .evaluate(genericRow(ONE_ROW));
+
+        // Then:
+        assertThat(result, is(ImmutableList.of("foo", "S1")));
+    }
+
+    @Test
+    public void shouldHandleCreateMap() {
+        // Given:
+        final Expression expression = analyzeQuery(
+            "SELECT MAP('foo' := 'foo', 'bar' := COL" + STRING_INDEX1 + ") FROM codegen_test EMIT CHANGES;", metaStore)
+            .getSelectExpressions()
+            .get(0)
+            .getExpression();
+
+        // When:
+        final Object result = codeGenRunner
+            .buildCodeGenFromParseTree(expression, "Map")
+            .evaluate(genericRow(ONE_ROW));
+
+        // Then:
+        assertThat(result, is(ImmutableMap.of("foo", "foo", "bar", "S1")));
+    }
+
     @Test
     public void shouldHandleInvalidJavaIdentifiers() {
         // Given:

diff --git a/ksql-engine/src/test/java/io/confluent/ksql/engine/InsertValuesExecutorTest.java b/ksql-engine/src/test/java/io/confluent/ksql/engine/InsertValuesExecutorTest.java
@@ -31,6 +31,7 @@
 import io.confluent.ksql.execution.ddl.commands.KsqlTopic;
 import io.confluent.ksql.execution.expression.tree.ArithmeticUnaryExpression;
 import io.confluent.ksql.execution.expression.tree.BooleanLiteral;
+import io.confluent.ksql.execution.expression.tree.CreateArrayExpression;
 import io.confluent.ksql.execution.expression.tree.DoubleLiteral;
 import io.confluent.ksql.execution.expression.tree.Expression;
 import io.confluent.ksql.execution.expression.tree.FunctionCall;
@@ -470,46 +471,48 @@ public void shouldHandleNegativeValueExpression() {
   @Test
   public void shouldHandleUdfs() {
     // Given:
-    givenSourceStreamWithSchema(SINGLE_ARRAY_SCHEMA, SerdeOption.none(), Optional.empty());
+    givenSourceStreamWithSchema(SINGLE_FIELD_SCHEMA, SerdeOption.none(), Optional.empty());
 
     final ConfiguredStatement<InsertValues> statement = givenInsertValuesStrings(
         ImmutableList.of("COL0"),
         ImmutableList.of(
             new FunctionCall(
-                FunctionName.of("AS_ARRAY"),
-                ImmutableList.of(new IntegerLiteral(1), new IntegerLiteral(2))))
+                FunctionName.of("SUBSTRING"),
+                ImmutableList.of(new StringLiteral("foo"), new IntegerLiteral(2))))
     );
 
     // When:
     executor.execute(statement, ImmutableMap.of(), engine, serviceContext);
 
     // Then:
-    verify(valueSerializer).serialize(TOPIC_NAME, new GenericRow(ImmutableList.of(ImmutableList.of(1, 2))));
+    verify(valueSerializer).serialize(TOPIC_NAME, new GenericRow(ImmutableList.of("oo")));
     verify(producer).send(new ProducerRecord<>(TOPIC_NAME, null, 1L, KEY, VALUE));
   }
 
   @Test
   public void shouldHandleNestedUdfs() {
     // Given:
-    givenSourceStreamWithSchema(SINGLE_MAP_SCHEMA, SerdeOption.none(), Optional.empty());
+    givenSourceStreamWithSchema(SINGLE_FIELD_SCHEMA, SerdeOption.none(), Optional.empty());
 
     final ConfiguredStatement<InsertValues> statement = givenInsertValuesStrings(
         ImmutableList.of("COL0"),
         ImmutableList.of(
             new FunctionCall(
-                FunctionName.of("AS_MAP"),
+                FunctionName.of("SUBSTRING"),
                 ImmutableList.of(
-                    new FunctionCall(FunctionName.of("AS_ARRAY"), ImmutableList.of(new StringLiteral("foo"))),
-                    new FunctionCall(FunctionName.of("AS_ARRAY"), ImmutableList.of(new IntegerLiteral(1)))
-                ))
-        )
+                    new FunctionCall(
+                        FunctionName.of("SUBSTRING"),
+                        ImmutableList.of(new StringLiteral("foo"), new IntegerLiteral(2))
+                    ),
+                    new IntegerLiteral(2))
+            ))
     );
 
     // When:
     executor.execute(statement, ImmutableMap.of(), engine, serviceContext);
 
     // Then:
-    verify(valueSerializer).serialize(TOPIC_NAME, new GenericRow(ImmutableList.of(ImmutableMap.of("foo", 1))));
+    verify(valueSerializer).serialize(TOPIC_NAME, new GenericRow(ImmutableList.of("o")));
     verify(producer).send(new ProducerRecord<>(TOPIC_NAME, null, 1L, KEY, VALUE));
   }
 

diff --git a/ksql-engine/src/test/java/io/confluent/ksql/engine/rewrite/ExpressionTreeRewriterTest.java b/ksql-engine/src/test/java/io/confluent/ksql/engine/rewrite/ExpressionTreeRewriterTest.java
@@ -26,6 +26,7 @@
 import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import io.confluent.ksql.engine.rewrite.ExpressionTreeRewriter.Context;
 import io.confluent.ksql.execution.expression.tree.ArithmeticBinaryExpression;
 import io.confluent.ksql.execution.expression.tree.ArithmeticUnaryExpression;
@@ -34,6 +35,8 @@
 import io.confluent.ksql.execution.expression.tree.Cast;
 import io.confluent.ksql.execution.expression.tree.ColumnReferenceExp;
 import io.confluent.ksql.execution.expression.tree.ComparisonExpression;
+import io.confluent.ksql.execution.expression.tree.CreateArrayExpression;
+import io.confluent.ksql.execution.expression.tree.CreateMapExpression;
 import io.confluent.ksql.execution.expression.tree.CreateStructExpression;
 import io.confluent.ksql.execution.expression.tree.CreateStructExpression.Field;
 import io.confluent.ksql.execution.expression.tree.DecimalLiteral;
@@ -532,6 +535,42 @@ public void shouldRewriteSubscriptExpression() {
     assertThat(rewritten, equalTo(new SubscriptExpression(parsed.getLocation(), expr1, expr2)));
   }
 
+  @Test
+  public void shouldRewriteCreateArrayExpression() {
+    // Given:
+    final CreateArrayExpression parsed = parseExpression("ARRAY['foo', col4[1]]");
+    final Expression firstVal = parsed.getValues().get(0);
+    final Expression secondVal = parsed.getValues().get(1);
+    when(processor.apply(firstVal, context)).thenReturn(expr1);
+    when(processor.apply(secondVal, context)).thenReturn(expr2);
+
+    // When:
+    final Expression rewritten = expressionRewriter.rewrite(parsed, context);
+
+    // Then:
+    assertThat(rewritten, equalTo(new CreateArrayExpression(ImmutableList.of(expr1, expr2))));
+  }
+
+  @Test
+  public void shouldRewriteCreateMapExpression() {
+    // Given:
+    final CreateMapExpression parsed = parseExpression("MAP('foo' := SUBSTRING('foo',0), 'bar' := col4[1])");
+    final Expression firstVal = parsed.getMap().get(new StringLiteral("foo"));
+    final Expression secondVal = parsed.getMap().get(new StringLiteral("bar"));
+    when(processor.apply(firstVal, context)).thenReturn(expr1);
+    when(processor.apply(secondVal, context)).thenReturn(expr2);
+    when(processor.apply(new StringLiteral("foo"), context)).thenReturn(new StringLiteral("foo"));
+    when(processor.apply(new StringLiteral("bar"), context)).thenReturn(new StringLiteral("bar"));
+
+    // When:
+    final Expression rewritten = expressionRewriter.rewrite(parsed, context);
+
+    // Then:
+    assertThat(rewritten,
+        equalTo(new CreateMapExpression(
+            ImmutableMap.of(new StringLiteral("foo"), expr1, new StringLiteral("bar"), expr2))));
+  }
+
   @Test
   public void shouldRewriteStructExpression() {
     // Given:

diff --git a/ksql-engine/src/test/java/io/confluent/ksql/function/udf/list/AsArrayTest.java b/ksql-engine/src/test/java/io/confluent/ksql/function/udf/list/AsArrayTest.java
diff --git a/ksql-execution/src/main/java/io/confluent/ksql/execution/codegen/CodeGenRunner.java b/ksql-execution/src/main/java/io/confluent/ksql/execution/codegen/CodeGenRunner.java
@@ -16,6 +16,8 @@
 package io.confluent.ksql.execution.codegen;
 
 import io.confluent.ksql.execution.expression.tree.ColumnReferenceExp;
+import io.confluent.ksql.execution.expression.tree.CreateArrayExpression;
+import io.confluent.ksql.execution.expression.tree.CreateMapExpression;
 import io.confluent.ksql.execution.expression.tree.CreateStructExpression;
 import io.confluent.ksql.execution.expression.tree.DereferenceExpression;
 import io.confluent.ksql.execution.expression.tree.Expression;
@@ -38,6 +40,7 @@
 import io.confluent.ksql.util.KsqlException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -185,6 +188,21 @@ public Void visitSubscriptExpression(final SubscriptExpression node, final Void
       return null;
     }
 
+    @Override
+    public Void visitCreateArrayExpression(final CreateArrayExpression exp, final Void context) {
+      exp.getValues().forEach(val -> process(val, context));
+      return null;
+    }
+
+    @Override
+    public Void visitCreateMapExpression(final CreateMapExpression exp, final Void context) {
+      for (Entry<Expression, Expression> entry : exp.getMap().entrySet()) {
+        process(entry.getKey(), context);
+        process(entry.getValue(), context);
+      }
+      return null;
+    }
+
     @Override
     public Void visitStructExpression(
         final CreateStructExpression exp,