[CALCITE-7608] Enhance Uncollect

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableUncollect.java

@@ -28,8 +28,7 @@
 import org.apache.calcite.runtime.SqlFunctions.FlatProductInputType;
 import org.apache.calcite.sql.type.MapSqlType;
 import org.apache.calcite.util.BuiltInMethod;
-
-import com.google.common.primitives.Ints;
+import org.apache.calcite.util.ImmutableBitSet;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -56,6 +55,19 @@
     assert getConvention() == child.getConvention();
   }
 
+  /** Creates an EnumerableUncollect with explicit control over which input
+   * fields are passed through unchanged and which are unnested.
+   *
+   * <p>Use {@link #create} unless you know what you're doing. */
+  public EnumerableUncollect(RelOptCluster cluster, RelTraitSet traitSet,
+      RelNode child, boolean withOrdinality, ImmutableBitSet passthroughFieldIndices,
+      ImmutableBitSet collectionFieldIndices, boolean outer, boolean expandStructFields) {
+    super(cluster, traitSet, child, withOrdinality, Collections.emptyList(),
+        passthroughFieldIndices, collectionFieldIndices, outer, expandStructFields);
+    assert getConvention() instanceof EnumerableConvention;
+    assert getConvention() == child.getConvention();
+  }
+
   /**
    * Creates an EnumerableUncollect.
    *
@@ -72,70 +84,98 @@
     return new EnumerableUncollect(cluster, traitSet, input, withOrdinality);
   }
 
+  /**
+   * Creates an EnumerableUncollect that unnests the collection-typed fields at
+   * {@code collectionFieldIndices}, passing through the fields at
+   * {@code passthroughFieldIndices} unchanged and dropping all others.
+   *
+   * @param traitSet                Trait set
+   * @param input                   Input relational expression
+   * @param withOrdinality          Whether output should contain an ORDINALITY column
+   * @param passthroughFieldIndices Fields from the input copied unchanged to the output
+   * @param collectionFieldIndices  Fields from the input which are unnested
+   * @param outer                   if true, preserves input rows with null/empty collections
+   *                                (LEFT JOIN); if false, drops them (INNER)
+   */
+  public static EnumerableUncollect create(RelTraitSet traitSet, RelNode input,
+      boolean withOrdinality, ImmutableBitSet passthroughFieldIndices,
+      ImmutableBitSet collectionFieldIndices, boolean outer, boolean expandStructFields) {
+    final RelOptCluster cluster = input.getCluster();
+    return new EnumerableUncollect(cluster, traitSet, input, withOrdinality,
+        passthroughFieldIndices, collectionFieldIndices, outer, expandStructFields);
+  }
+
   @Override public EnumerableUncollect copy(RelTraitSet traitSet,
       RelNode newInput) {
-    return new EnumerableUncollect(getCluster(), traitSet, newInput,
-        withOrdinality);
+    return new EnumerableUncollect(getCluster(), traitSet, newInput, withOrdinality,
+        getPassthroughFieldIndices(), getCollectionFieldIndices(), isOuter,
+        expandStructFields);
   }
 
+  /** Implements Uncollect: a subset of input fields (possibly all of them)
+   * are unnested, the rest are either passed through unchanged or dropped,
+   * and (if {@link #isOuter}) rows with null/empty collections are preserved
+   * with NULL-padded element columns. */
   @Override public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
+    final ImmutableBitSet passthroughIndices = getPassthroughFieldIndices();
+    final ImmutableBitSet collIndices = getCollectionFieldIndices();
     final BlockBuilder builder = new BlockBuilder();
     final EnumerableRel child = (EnumerableRel) getInput();
-    final Result result = implementor.visitChild(this, 0, child, pref);
+    final Result childResult = implementor.visitChild(this, 0, child, pref);
     final PhysType physType =
-        PhysTypeImpl.of(
-            implementor.getTypeFactory(),
-            getRowType(),
-            JavaRowFormat.LIST);
-
-    // final Enumerable<List<Employee>> child = <<child adapter>>;
-    // return child.selectMany(FLAT_ZIP);
-    final Expression child_ =
-        builder.append(
-            "child", result.block);
+        PhysTypeImpl.of(implementor.getTypeFactory(), getRowType(), JavaRowFormat.LIST);
 
+    // Compute element width and kind for each collection field.
+    final List<RelDataTypeField> inputFields =
+        getInput().getRowType().getFieldList();
     final List<Integer> fieldCounts = new ArrayList<>();
     final List<FlatProductInputType> inputTypes = new ArrayList<>();
-
-    Expression lambdaForStructWithSingleItem = null;
-    for (RelDataTypeField field : child.getRowType().getFieldList()) {
-      final RelDataType type = field.getType();
-      if (type instanceof MapSqlType) {
+    for (int idx : collIndices) {
+      final RelDataType ft = inputFields.get(idx).getType();
+      if (ft instanceof MapSqlType) {
         fieldCounts.add(2);
         inputTypes.add(FlatProductInputType.MAP);
       } else {
-        final RelDataType elementType = getComponentTypeOrThrow(type);
-        if (elementType.isStruct()) {
-          if (elementType.getFieldCount() == 1 && child.getRowType().getFieldList().size() == 1
-              && !withOrdinality) {
-            // Solves CALCITE-4063: if we are processing a single field, which is a struct with a
-            // single item inside, and no ordinality; the result must be a scalar, hence use a
-            // special lambda that does not return lists, but the (single) items within those lists
-            lambdaForStructWithSingleItem = Expressions.call(BuiltInMethod.FLAT_LIST.method);
-          } else {
-            fieldCounts.add(elementType.getFieldCount());
-            inputTypes.add(FlatProductInputType.LIST);
-          }
+        final RelDataType ct = getComponentTypeOrThrow(ft);
+        if (ct.isStruct() && expandStructFields) {
+          fieldCounts.add(ct.getFieldCount());
+          inputTypes.add(FlatProductInputType.LIST);
+        } else if (ct.isStruct()) {
+          // A struct element kept whole occupies a single output column,
+          // like a scalar element, but its row value must be converted from
+          // the collection's internal list representation to Object[].
+          fieldCounts.add(-1);
+          inputTypes.add(FlatProductInputType.STRUCT);
         } else {
+          // Scalar elements occupy a single output column, unconverted.
           fieldCounts.add(-1);
           inputTypes.add(FlatProductInputType.SCALAR);
         }
       }
     }
 
-    final Expression lambda = lambdaForStructWithSingleItem != null
-        ? lambdaForStructWithSingleItem
-        : Expressions.call(BuiltInMethod.FLAT_ZIP.method,
-            Expressions.constant(Ints.toArray(fieldCounts)),
-            Expressions.constant(withOrdinality),
+    final Expression childExpression = builder.append("child", childResult.block);
+
+    // SqlFunctions.flatUncollect(passthroughIndices, collectionIndices,
+    //     fieldCounts, inputTypes, withOrdinality, inputFieldCount, outer)
+    final Expression flatUncollectFn =
+        Expressions.call(BuiltInMethod.FLAT_UNCOLLECT.method,
+            Expressions.constant(passthroughIndices.toArray()),
+            Expressions.constant(collIndices.toArray()),
+            // Following converts from a List<Integer> to int[] (not Integer[]).
+            Expressions.constant(fieldCounts.stream().mapToInt(i -> i).toArray()),
             Expressions.constant(
-                inputTypes.toArray(new FlatProductInputType[0])));
+                inputTypes.toArray(new FlatProductInputType[0])),
+            Expressions.constant(withOrdinality),
+            Expressions.constant(inputFields.size()),
+            Expressions.constant(isOuter));
+
     builder.add(
         Expressions.return_(null,
-            Expressions.call(child_,
+            Expressions.call(childExpression,
                 BuiltInMethod.SELECT_MANY.method,
-                lambda)));
+                flatUncollectFn)));
+
     return implementor.result(physType, builder.toBlock());
   }
-
 }

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableUncollectRule.java

@@ -49,6 +49,8 @@ protected EnumerableUncollectRule(Config config) {
         convert(input,
             input.getTraitSet().replace(EnumerableConvention.INSTANCE));
     return EnumerableUncollect.create(traitSet, newInput,
-        uncollect.withOrdinality);
+        uncollect.withOrdinality, uncollect.getPassthroughFieldIndices(),
+        uncollect.getCollectionFieldIndices(), uncollect.isOuter,
+        uncollect.expandStructFields);
   }
 }

core/src/main/java/org/apache/calcite/interpreter/UncollectNode.java

@@ -17,11 +17,20 @@
 package org.apache.calcite.interpreter;
 
 import org.apache.calcite.rel.core.Uncollect;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.calcite.sql.type.SqlTypeName;
 
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 
+import static java.util.Objects.requireNonNull;
+
 /**
  * Interpreter node that implements a
  * {@link org.apache.calcite.rel.core.Uncollect}.
@@ -33,6 +42,17 @@
   }
 
   @Override public void run() throws InterruptedException {
+    final int inputFieldCount = rel.getInput().getRowType().getFieldCount();
+    if (rel.getPassthroughFieldIndices().isEmpty()
+        && rel.getCollectionFieldIndices().cardinality() == inputFieldCount) {
+      runPlain();
+    } else {
+      runGeneralized();
+    }
+  }
+
+  /** Unnests every input field; every input field must be a collection. */
+  private void runPlain() throws InterruptedException {
     Row row = null;
     while ((row = source.receive()) != null) {
       for (Object value : row.getValues()) {
@@ -67,4 +87,116 @@
       }
     }
   }
+
+  /**
+   * Unnests the collection fields and prepends the passthrough fields,
+   * mirroring {@code SqlFunctions.flatUncollect}: multiple collections zip
+   * with {@code null} padding, a {@code null} collection counts as empty,
+   * and when all collections are empty {@code isOuter} decides between one
+   * {@code null}-padded row and dropping the input row.
+   */
+  private void runGeneralized() throws InterruptedException {
+    final List<RelDataTypeField> inputFields = rel.getInput().getRowType().getFieldList();
+    final List<Integer> passthroughs = rel.getPassthroughFieldIndices().toList();
+    final List<Integer> collections = rel.getCollectionFieldIndices().toList();
+
+    // Output width of each collection: 2 for maps (key, value), the component
+    // field count for struct elements, otherwise 1.
+    final int[] widths = new int[collections.size()];
+    int elementWidth = 0;
+    for (int c = 0; c < collections.size(); c++) {
+      final RelDataType type = inputFields.get(collections.get(c)).getType();
+      if (type.getSqlTypeName() == SqlTypeName.MAP) {
+        widths[c] = 2;
+      } else {
+        final RelDataType component =
+            requireNonNull(type.getComponentType(), "componentType");
+        widths[c] = component.isStruct() ? component.getFieldCount() : 1;
+      }
+      elementWidth += widths[c];
+    }
+    final int outputWidth =
+        passthroughs.size() + elementWidth + (rel.withOrdinality ? 1 : 0);
+
+    Row row = null;
+    while ((row = source.receive()) != null) {
+      final @Nullable Object[] values = row.getValues();
+      final List<List<Object>> elements = new ArrayList<>();
+      int maxSize = 0;
+      for (int index : collections) {
+        final Object value = values[index];
+        final List<Object> list;
+        if (value == null) {
+          list = Collections.emptyList();
+        } else if (value instanceof List) {
+          list = (List<Object>) value;
+        } else if (value instanceof Map) {
+          list = new ArrayList<>(((Map<Object, Object>) value).entrySet());
+        } else {
+          throw new UnsupportedOperationException(
+              String.format(Locale.ROOT,
+                  "Invalid type: %s for unnest.",
+                  value.getClass().getCanonicalName()));
+        }
+        elements.add(list);
+        maxSize = Math.max(maxSize, list.size());
+      }
+
+      if (maxSize == 0) {
+        if (rel.isOuter) {
+          final @Nullable Object[] out = new Object[outputWidth];
+          int outCol = 0;
+          for (int p : passthroughs) {
+            out[outCol++] = values[p];
+          }
+          sink.send(Row.of(out));
+        }
+        continue;
+      }
+
+      for (int i = 0; i < maxSize; i++) {
+        final @Nullable Object[] out = new Object[outputWidth];
+        int outCol = 0;
+        for (int p : passthroughs) {
+          out[outCol++] = values[p];
+        }
+        for (int c = 0; c < collections.size(); c++) {
+          final List<Object> list = elements.get(c);
+          final Object element = i < list.size() ? list.get(i) : null;
+          outCol = appendElement(out, outCol, widths[c], element);
+        }
+        if (rel.withOrdinality) {
+          out[outCol] = i + 1;
+        }
+        sink.send(Row.of(out));
+      }
+    }
+  }
+
+  /** Writes one collection element into {@code out} starting at {@code outCol},
+   * expanding map entries and struct elements into their columns; a {@code null}
+   * element leaves its columns {@code null}. Returns the next free column. */
+  private static int appendElement(@Nullable Object [] out, int outCol, int width,
+      @Nullable Object element) {
+    if (element instanceof Map.Entry) {
+      final Map.Entry<?, ?> entry = (Map.Entry<?, ?>) element;
+      out[outCol] = entry.getKey();
+      out[outCol + 1] = entry.getValue();
+    } else if (width == 1 || element == null) {
+      out[outCol] = element;
+    } else if (element instanceof Object[]) {
+      System.arraycopy((Object[]) element, 0, out, outCol, width);
+    } else if (element instanceof List) {
+      final List<?> struct = (List<?>) element;
+      for (int i = 0; i < width; i++) {
+        out[outCol + i] = struct.get(i);
+      }
+    } else {
+      throw new UnsupportedOperationException(
+          String.format(Locale.ROOT,
+              "Invalid element type: %s for unnest.",
+              element.getClass().getCanonicalName()));
+    }
+    return outCol + width;
+  }
 }