wasabii commented on code in PR #4922:
URL: https://github.com/apache/calcite/pull/4922#discussion_r3300327096


##########
core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableTableModify.java:
##########
@@ -92,88 +97,304 @@ public EnumerableTableModify(RelOptCluster cluster, 
RelTraitSet traits,
         ModifiableTable.class.isAssignableFrom(
             Types.toClass(expression.getType())),
         "not assignable from type %s", expression.getType());
+
     builder.add(
         Expressions.declare(
             Modifier.FINAL,
             collectionParameter,
             Expressions.call(
                 expression,
-                BuiltInMethod.MODIFIABLE_TABLE_GET_MODIFIABLE_COLLECTION
-                    .method)));
+                
BuiltInMethod.MODIFIABLE_TABLE_GET_MODIFIABLE_COLLECTION.method)));
+
+    final PhysType physType =
+        PhysTypeImpl.of(
+            implementor.getTypeFactory(),
+            getRowType(),
+            pref == Prefer.ARRAY ? JavaRowFormat.ARRAY : JavaRowFormat.SCALAR);
+
+    switch (getOperation()) {
+    case INSERT:
+      return implementInsert(implementor, result, builder, childExp, 
collectionParameter, physType);
+    case UPDATE:
+      return implementUpdate(implementor, builder, childExp, 
collectionParameter, physType);
+    case DELETE:
+      return implementDelete(implementor, builder, childExp, 
collectionParameter, physType);
+    default:
+      throw new AssertionError("unsupported operation: " + getOperation());
+    }
+  }
+
+  /** Generates code for an UPDATE statement.
+   *
+   * <p>The child produces, for each row matched by the WHERE clause, a row of
+   * {@code tableFieldCount + M} fields:
+   * {@code [originalField_0, ..., originalField_N-1, newValue_0, ..., 
newValue_M-1]}.
+   * The first {@code N} fields are the <em>entire</em> original table row (all
+   * columns, not just those being updated); the trailing {@code M =
+   * updateColumnList.size()} fields are the new values, one per column named 
in
+   * the SET clause. Filtering by WHERE has already been applied upstream, so
+   * every source row corresponds to an existing row in the modifiable 
collection
+   * and can be located by full-row content equality.
+   */
+  private Result implementUpdate(
+      EnumerableRelImplementor implementor,
+      BlockBuilder builder,
+      Expression childExp,
+      ParameterExpression collectionParameter,
+      PhysType physType) {
+    final List<String> updateCols = requireNonNull(getUpdateColumnList());
+    final List<RelDataTypeField> tableFields = 
table.getRowType().getFieldList();
+    final int tableFieldCount = tableFields.size();
+    final int[] updateColumnIndices = new int[updateCols.size()];
+    for (int i = 0; i < updateCols.size(); i++) {
+      final String colName = updateCols.get(i);
+      int found = -1;
+      for (int j = 0; j < tableFields.size(); j++) {
+        if (tableFields.get(j).getName().equals(colName)) {
+          found = j;
+          break;
+        }
+      }
+      if (found < 0) {
+        throw new AssertionError("column '" + colName + "' not found in 
table");
+      }
+      updateColumnIndices[i] = found;
+    }
+
+    // Build the three lambdas required by ExtendedEnumerable.update:
+    //   sinkKeySelector:   row -> Arrays.asList(row)
+    //   sourceKeySelector: row -> Arrays.asList(Arrays.copyOf(row, N))
+    //   sourceTransform:   row -> applyUpdate(row, N, updateColumnIndices)
+    final ParameterExpression sinkRow =
+        Expressions.parameter(Object[].class, "sinkRow");
+    final Expression sinkKeySelector =
+        Expressions.lambda(Function1.class,
+            Expressions.call(Arrays.class, "asList", sinkRow),
+            sinkRow);
+
+    final ParameterExpression srcKeyRow =
+        Expressions.parameter(Object[].class, "row");
+    final Expression sourceKeySelector =
+        Expressions.lambda(Function1.class,
+            Expressions.call(Arrays.class, "asList",
+                Expressions.call(Arrays.class, "copyOf",
+                    srcKeyRow, Expressions.constant(tableFieldCount))),
+            srcKeyRow);
+
+    final ParameterExpression srcXformRow =
+        Expressions.parameter(Object[].class, "row");
+    final Expression sourceTransform =
+        Expressions.lambda(Function1.class,
+            Expressions.call(EnumerableTableModify.class, "applyUpdate",
+                srcXformRow,
+                Expressions.constant(tableFieldCount),
+                Expressions.constant(updateColumnIndices)),
+            srcXformRow);
+
+    final Expression updateCountExp =
+        builder.append(
+            "updateCount",
+            Expressions.call(
+                childExp,
+                BuiltInMethod.UPDATE.method,
+                Expressions.convert_(collectionParameter, List.class),
+                sinkKeySelector,
+                sourceKeySelector,
+                sourceTransform));
+
+    builder.add(
+        Expressions.return_(
+            null,
+            Expressions.call(
+                BuiltInMethod.SINGLETON_ENUMERABLE.method,
+                Expressions.convert_(updateCountExp, long.class))));
+
+    return implementor.result(physType, builder.toBlock());
+  }
+
+  /** Generates code for a DELETE statement.
+   *
+   * <p>{@code Object[].equals()} is reference equality, so
+   * {@link Collection#removeAll} would fail to match rows that were copied
+   * during query planning or execution. {@link #applyDelete} uses
+   * {@link Arrays#asList} wrapping to get value-based equality for
+   * {@code Object[]} rows, so it works regardless of whether the child yields
+   * the same {@code Object[]} references that live in the backing collection.
+   */
+  private Result implementDelete(
+      EnumerableRelImplementor implementor,
+      BlockBuilder builder,
+      Expression childExp,
+      ParameterExpression collectionParameter,
+      PhysType physType) {
+
+    final Expression countParameter =
+        builder.append(
+            "count",
+            Expressions.call(collectionParameter, "size"),
+            false);
+
+    builder.add(
+        Expressions.statement(
+            Expressions.call(
+                EnumerableTableModify.class, "applyDelete",
+                childExp, collectionParameter)));
+
+    final Expression updatedCountParameter =
+        builder.append(
+            "updatedCount",
+            Expressions.call(collectionParameter, "size"),
+            false);
+
+    builder.add(
+        Expressions.return_(
+            null,
+            Expressions.call(
+                BuiltInMethod.SINGLETON_ENUMERABLE.method,
+                Expressions.convert_(
+                    Expressions.subtract(countParameter, 
updatedCountParameter),
+                    long.class))));
+
+    return implementor.result(physType, builder.toBlock());
+  }
+
+  /** Generates code for an INSERT statement.
+   *
+   * <p>When the child row type differs from the table row type (the common 
case
+   * for DML, since the node's output is {@code {ROWCOUNT: BIGINT}}), the child
+   * expression is wrapped in a type-conversion projection before being 
streamed
+   * into the backing collection.
+   */
+  private Result implementInsert(
+      EnumerableRelImplementor implementor,
+      Result result,
+      BlockBuilder builder,
+      Expression childExp,
+      ParameterExpression collectionParameter,
+      PhysType physType) {
     final Expression countParameter =
         builder.append(
             "count",
             Expressions.call(collectionParameter, "size"),
             false);
+
     Expression convertedChildExp;
     if (!getInput().getRowType().equals(getRowType())) {
       final JavaTypeFactory typeFactory =
           (JavaTypeFactory) getCluster().getTypeFactory();
       final JavaRowFormat format = EnumerableTableScan.deduceFormat(table);
-      PhysType physType =
+      PhysType tablePhysType =
           PhysTypeImpl.of(typeFactory, table.getRowType(), format);
       List<Expression> expressionList = new ArrayList<>();
       final PhysType childPhysType = result.physType;
       final ParameterExpression o_ =
           Expressions.parameter(childPhysType.getJavaRowType(), "o");
-      final int fieldCount =
-          childPhysType.getRowType().getFieldCount();
+      final int fieldCount = childPhysType.getRowType().getFieldCount();
       for (int i = 0; i < fieldCount; i++) {
         expressionList.add(
-            childPhysType.fieldReference(o_, i, physType.getJavaFieldType(i)));
+            childPhysType.fieldReference(o_, i, 
tablePhysType.getJavaFieldType(i)));
       }
+
       convertedChildExp =
           builder.append(
               "convertedChild",
               Expressions.call(
                   childExp,
                   BuiltInMethod.SELECT.method,
-                  Expressions.lambda(
-                      physType.record(expressionList), o_)));
+                  Expressions.lambda(tablePhysType.record(expressionList), 
o_)));
     } else {
       convertedChildExp = childExp;
     }
-    final Method method;
-    switch (getOperation()) {
-    case INSERT:
-      method = BuiltInMethod.INTO.method;
-      break;
-    case DELETE:
-      method = BuiltInMethod.REMOVE_ALL.method;
-      break;
-    default:
-      throw new AssertionError(getOperation());
-    }
+
     builder.add(
         Expressions.statement(
             Expressions.call(
-                convertedChildExp, method, collectionParameter)));
+                convertedChildExp, BuiltInMethod.INTO.method, 
collectionParameter)));
+
     final Expression updatedCountParameter =
         builder.append(
             "updatedCount",
             Expressions.call(collectionParameter, "size"),
             false);
+
     builder.add(
         Expressions.return_(
             null,
             Expressions.call(
                 BuiltInMethod.SINGLETON_ENUMERABLE.method,
                 Expressions.convert_(
-                    Expressions.condition(
-                        Expressions.greaterThanOrEqual(
-                            updatedCountParameter, countParameter),
-                        Expressions.subtract(
-                            updatedCountParameter, countParameter),
-                        Expressions.subtract(
-                            countParameter, updatedCountParameter)),
+                    Expressions.subtract(updatedCountParameter, 
countParameter),

Review Comment:
   If you mean can it ever be negative? No. This function only inserts.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to