This is an automated email from the ASF dual-hosted git repository.

jhyde pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/calcite.git

commit 9c741e4415362c3358acea09e223d64d9b36b1ae
Author: Julian Hyde <[email protected]>
AuthorDate: Wed Oct 13 17:18:42 2021 -0700

    [CALCITE-4719] Add variants of RexSubQuery that collect sub-queries into 
MULTISET, ARRAY and MAP collections
    
    In RexSubQuery, add methods `array`, `map`, `multiset`.
    
    In RelBuilder, add methods `arrayQuery`, `mapQuery`,
    `multisetQuery`.
    
    Close apache/calcite#2582
---
 .../java/org/apache/calcite/rex/RexSubQuery.java   | 31 ++++++++-
 .../java/org/apache/calcite/tools/RelBuilder.java  | 73 ++++++++++++++++++++++
 .../org/apache/calcite/test/RelBuilderTest.java    | 62 ++++++++++++++++++
 site/_docs/algebra.md                              |  6 ++
 4 files changed, 171 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/apache/calcite/rex/RexSubQuery.java 
b/core/src/main/java/org/apache/calcite/rex/RexSubQuery.java
index 076c6ba..e1db6f8 100644
--- a/core/src/main/java/org/apache/calcite/rex/RexSubQuery.java
+++ b/core/src/main/java/org/apache/calcite/rex/RexSubQuery.java
@@ -108,7 +108,9 @@ public class RexSubQuery extends RexCall {
   /** Creates a scalar sub-query. */
   public static RexSubQuery scalar(RelNode rel) {
     final List<RelDataTypeField> fieldList = rel.getRowType().getFieldList();
-    assert fieldList.size() == 1;
+    if (fieldList.size() != 1) {
+      throw new IllegalArgumentException();
+    }
     final RelDataTypeFactory typeFactory = rel.getCluster().getTypeFactory();
     final RelDataType type =
         typeFactory.createTypeWithNullability(fieldList.get(0).getType(), 
true);
@@ -116,6 +118,33 @@ public class RexSubQuery extends RexCall {
         ImmutableList.of(), rel);
   }
 
+  /** Creates an ARRAY sub-query. */
+  public static RexSubQuery array(RelNode rel) {
+    final RelDataTypeFactory typeFactory = rel.getCluster().getTypeFactory();
+    final RelDataType type =
+        typeFactory.createArrayType(rel.getRowType(), -1L);
+    return new RexSubQuery(type, SqlStdOperatorTable.ARRAY_QUERY,
+        ImmutableList.of(), rel);
+  }
+
+  /** Creates a MULTISET sub-query. */
+  public static RexSubQuery multiset(RelNode rel) {
+    final RelDataTypeFactory typeFactory = rel.getCluster().getTypeFactory();
+    final RelDataType type =
+        typeFactory.createMultisetType(rel.getRowType(), -1L);
+    return new RexSubQuery(type, SqlStdOperatorTable.MULTISET_QUERY,
+        ImmutableList.of(), rel);
+  }
+
+  /** Creates a MAP sub-query. */
+  public static RexSubQuery map(RelNode rel) {
+    final RelDataTypeFactory typeFactory = rel.getCluster().getTypeFactory();
+    final RelDataType type =
+        typeFactory.createMultisetType(rel.getRowType(), -1L);
+    return new RexSubQuery(type, SqlStdOperatorTable.MAP_QUERY,
+        ImmutableList.of(), rel);
+  }
+
   @Override public <R> R accept(RexVisitor<R> visitor) {
     return visitor.visitSubQuery(this);
   }
diff --git a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java 
b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java
index 5132636..df81877 100644
--- a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java
+++ b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java
@@ -928,6 +928,79 @@ public class RelBuilder {
     return RexSubQuery.scalar(f.apply(this));
   }
 
+  /** Creates an ARRAY sub-query.
+   *
+   * <p>For example,
+   *
+   * <pre>{@code
+   * b.scan("Depts")
+   *     .project(
+   *         b.field("deptno")
+   *         b.arrayQuery(b2 ->
+   *             b2.scan("Emps")
+   *                 .build()))
+   * }</pre>
+   *
+   * <p>is equivalent to the SQL
+   *
+   * <pre>{@code
+   * SELECT deptno, ARRAY (SELECT * FROM Emps)
+   * FROM Depts
+   * }</pre> */
+  @Experimental
+  public RexSubQuery arrayQuery(Function<RelBuilder, RelNode> f) {
+    return RexSubQuery.array(f.apply(this));
+  }
+
+  /** Creates a MULTISET sub-query.
+   *
+   * <p>For example,
+   *
+   * <pre>{@code
+   * b.scan("Depts")
+   *     .project(
+   *         b.field("deptno")
+   *         b.multisetQuery(b2 ->
+   *             b2.scan("Emps")
+   *                 .build()))
+   * }</pre>
+   *
+   * <p>is equivalent to the SQL
+   *
+   * <pre>{@code
+   * SELECT deptno, MULTISET (SELECT * FROM Emps)
+   * FROM Depts
+   * }</pre> */
+  @Experimental
+  public RexSubQuery multisetQuery(Function<RelBuilder, RelNode> f) {
+    return RexSubQuery.multiset(f.apply(this));
+  }
+
+  /** Creates a MAP sub-query.
+   *
+   * <p>For example,
+   *
+   * <pre>{@code
+   * b.scan("Depts")
+   *     .project(
+   *         b.field("deptno")
+   *         b.multisetQuery(b2 ->
+   *             b2.scan("Emps")
+   *                 .project(b2.field("empno"), b2.field("job"))
+   *                 .build()))
+   * }</pre>
+   *
+   * <p>is equivalent to the SQL
+   *
+   * <pre>{@code
+   * SELECT deptno, MAP (SELECT empno, job FROM Emps)
+   * FROM Depts
+   * }</pre> */
+  @Experimental
+  public RexSubQuery mapQuery(Function<RelBuilder, RelNode> f) {
+    return RexSubQuery.map(f.apply(this));
+  }
+
   /** Creates an AND. */
   public RexNode and(RexNode... operands) {
     return and(ImmutableList.copyOf(operands));
diff --git a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java 
b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java
index bdfe0c7..c210874 100644
--- a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java
@@ -2609,6 +2609,68 @@ public class RelBuilderTest {
     assertThat(f.apply(createBuilder()), hasTree(expected));
   }
 
+  @Test void testArrayQuery() {
+    // Equivalent SQL:
+    //   SELECT deptno, ARRAY (SELECT * FROM Emp)
+    //   FROM Dept AS d
+    final Function<RelBuilder, RelNode> f = b ->
+        b.scan("DEPT")
+            .project(
+                b.field("DEPTNO"),
+                b.arrayQuery(b2 ->
+                        b2.scan("EMP")
+                            .build()))
+            .build();
+
+    final String expected = "LogicalProject(DEPTNO=[$0], $f1=[ARRAY({\n"
+        + "LogicalTableScan(table=[[scott, EMP]])\n"
+        + "})])\n"
+        + "  LogicalTableScan(table=[[scott, DEPT]])\n";
+    assertThat(f.apply(createBuilder()), hasTree(expected));
+  }
+
+  @Test void testMultisetQuery() {
+    // Equivalent SQL:
+    //   SELECT deptno, MULTISET (SELECT * FROM Emp)
+    //   FROM Dept AS d
+    final Function<RelBuilder, RelNode> f = b ->
+        b.scan("DEPT")
+            .project(
+                b.field("DEPTNO"),
+                b.multisetQuery(b2 ->
+                        b2.scan("EMP")
+                            .build()))
+            .build();
+
+    final String expected = "LogicalProject(DEPTNO=[$0], $f1=[MULTISET({\n"
+        + "LogicalTableScan(table=[[scott, EMP]])\n"
+        + "})])\n"
+        + "  LogicalTableScan(table=[[scott, DEPT]])\n";
+    assertThat(f.apply(createBuilder()), hasTree(expected));
+  }
+
+  @Test void testMapQuery() {
+    // Equivalent SQL:
+    //   SELECT deptno, MAP (SELECT empno, job FROM Emp)
+    //   FROM Dept AS d
+    final Function<RelBuilder, RelNode> f = b ->
+        b.scan("DEPT")
+            .project(
+                b.field("DEPTNO"),
+                b.mapQuery(b2 ->
+                        b2.scan("EMP")
+                            .project(b2.field("EMPNO"), b2.field("JOB"))
+                            .build()))
+            .build();
+
+    final String expected = "LogicalProject(DEPTNO=[$0], $f1=[MAP({\n"
+        + "LogicalProject(EMPNO=[$0], JOB=[$2])\n"
+        + "  LogicalTableScan(table=[[scott, EMP]])\n"
+        + "})])\n"
+        + "  LogicalTableScan(table=[[scott, DEPT]])\n";
+    assertThat(f.apply(createBuilder()), hasTree(expected));
+  }
+
   @Test void testAlias() {
     // Equivalent SQL:
     //   SELECT *
diff --git a/site/_docs/algebra.md b/site/_docs/algebra.md
index 7a97140..66b2792 100644
--- a/site/_docs/algebra.md
+++ b/site/_docs/algebra.md
@@ -448,6 +448,9 @@ added to the stack.
 The following methods convert a sub-query into a scalar value (a `BOOLEAN` in
 the case of `in`, `exists`, `some`, `all`, `unique`;
 any scalar type for `scalarQuery`).
+an `ARRAY` for `arrayQuery`,
+a `MAP` for `mapQuery`,
+and a `MULTISET` for `multisetQuery`).
 
 In all the following, `relFn` is a function that takes a `RelBuilder` argument
 and returns a `RelNode`. You typically implement it as a lambda; the method
@@ -457,8 +460,11 @@ returns the `RelNode` that is to be the sub-query.
 | Method              | Description
 |:------------------- |:-----------
 | `all(expr, op, relFn)` | Returns whether *expr* has a particular relation to 
all of the values of the sub-query
+| `arrayQuery(relFn)` | Returns the rows of a sub-query as an `ARRAY`
 | `exists(relFn)` | Tests whether sub-query is non-empty
 | `in(expr, relFn)`<br/>`in(exprList, relFn)` | Tests whether a value occurs 
in a sub-query
+| `mapQuery(relFn)` | Returns the rows of a sub-query as a `MAP`
+| `multisetQuery(relFn)` | Returns the rows of a sub-query as a `MULTISET`
 | `scalarQuery(relFn)` | Returns the value of the sole column of the sole row 
of a sub-query
 | `some(expr, op, relFn)` | Returns whether *expr* has a particular relation 
to one or more of the values of the sub-query
 | `unique(relFn)` | Returns whether the rows of a sub-query are unique

Reply via email to