Add SINGLE_VALUE aggregate function

sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/transform/BeamBuiltinAggregations.java

@@ -60,7 +60,17 @@ public class BeamBuiltinAggregations {
   public static final Map<String, Function<Schema.FieldType, CombineFn<?, ?, ?>>>
       BUILTIN_AGGREGATOR_FACTORIES =
           ImmutableMap.<String, Function<Schema.FieldType, CombineFn<?, ?, ?>>>builder()
+              // SINGLE_VALUE is emitted by Calcite to enforce the cardinality of a scalar
+              // subquery (a subquery used as a scalar must yield exactly one row). The single
+              // input value is returned as-is; unlike COUNT/SUM it must not drop nulls, so a
+              // scalar subquery evaluating to NULL surfaces NULL.
+              .put("SINGLE_VALUE", typeName -> new SingleValue<>())
               .put("ANY_VALUE", typeName -> Sample.anyValueCombineFn())
+              // SINGLE_VALUE is emitted by Calcite to enforce the cardinality of a scalar
+              // subquery (a subquery used as a scalar must yield exactly one row). The single
+              // input value is returned as-is; unlike COUNT/SUM it must not drop nulls, so a
+              // scalar subquery evaluating to NULL surfaces NULL.
+              .put("SINGLE_VALUE", typeName -> Sample.anyValueCombineFn())
               // Drop null elements for these aggregations.
               .put("COUNT", typeName -> new DropNullFnWithDefault(Count.combineFn()))
               .put("MAX", typeName -> new DropNullFn(BeamBuiltinAggregations.createMax(typeName)))
@@ -693,4 +703,52 @@ public Long extractOutput(BitXOr.Accum accumulator) {
       return accumulator.bitXOr;
     }
   }
+
+  /**
+   * {@link CombineFn} for SINGLE_VALUE to enforce that a scalar subquery returns exactly one row.
+   */
+  public static class SingleValue<T> extends CombineFn<T, java.util.List<T>, T> {
+    @Override
+    public java.util.List<T> createAccumulator() {
+      return new java.util.ArrayList<>();
+    }
+
+    @Override
+    public java.util.List<T> addInput(java.util.List<T> accumulator, T input) {
+      if (accumulator.size() < 2) {
+        accumulator.add(input);
+      }
+      return accumulator;
+    }
+
+    @Override
+    public java.util.List<T> mergeAccumulators(Iterable<java.util.List<T>> accumulators) {
+      java.util.List<T> merged = createAccumulator();
+      for (java.util.List<T> accum : accumulators) {
+        merged.addAll(accum);
+        if (merged.size() >= 2) {
+          merged = new java.util.ArrayList<>(merged.subList(0, 2));
+        }
+      }
+      return merged;
+    }
+
+    @Override
+    public T extractOutput(java.util.List<T> accumulator) {
+      if (accumulator.isEmpty()) {
+        return null;
+      }
+      if (accumulator.size() > 1) {
+        throw new IllegalArgumentException("Subquery returned more than one row");
+      }
+      return accumulator.get(0);
+    }
+
+    @Override
+    public Coder<java.util.List<T>> getAccumulatorCoder(
+        CoderRegistry registry, Coder<T> inputCoder) {
+      return org.apache.beam.sdk.coders.ListCoder.of(
+          org.apache.beam.sdk.coders.NullableCoder.of(inputCoder));
+    }
+  }
 }

sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlDslAggregationTest.java

@@ -373,6 +373,53 @@
     pipeline.run().waitUntilFinish();
   }
 
+  /** GROUP-BY with the SINGLE_VALUE aggregation function. */
+  @Test
+  public void testSingleValueFunction() throws Exception {
+    pipeline.enableAbandonedNodeEnforcement(false);
+
+    Schema schema = Schema.builder().addInt32Field("key").addInt32Field("col").build();
+
+    PCollection<Row> inputRows =
+        pipeline
+            .apply(
+                Create.of(
+                    TestUtils.rowsBuilderOf(schema)
+                        .addRows(
+                            0, 1,
+                            0, 2,
+                            1, 3,
+                            2, 4,
+                            2, 5)
+                        .getRows()))
+            .setRowSchema(schema);
+
+    String sql = "SELECT key, SINGLE_VALUE(col) as single_value FROM PCOLLECTION GROUP BY key";
+
+    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));
+
+    Map<Integer, List<Integer>> allowedTuples = new HashMap<>();
+    allowedTuples.put(0, Arrays.asList(1, 2));
+    allowedTuples.put(1, Arrays.asList(3));
+    allowedTuples.put(2, Arrays.asList(4, 5));
+
+    PAssert.that(result)
+        .satisfies(
+            input -> {
+              Iterator<Row> iter = input.iterator();
+              while (iter.hasNext()) {
+                Row row = iter.next();
+                List<Integer> values = allowedTuples.remove(row.getInt32("key"));
+                assertTrue(values != null);
+                assertTrue(values.contains(row.getInt32("single_value")));
+              }
+              assertTrue(allowedTuples.isEmpty());
+              return null;
+            });
+
+    pipeline.run();
+  }
+
   /** GROUP-BY with the any_value aggregation function. */
   @Test
   public void testAnyValueFunction() throws Exception {
@@ -420,6 +467,78 @@
     pipeline.run();
   }
 
+  /** GROUP-BY with the SINGLE_VALUE aggregation function. */
+  @Test
+  public void testSingleValueFunction() throws Exception {
+    pipeline.enableAbandonedNodeEnforcement(false);
+
+    Schema schema = Schema.builder().addInt32Field("key").addInt32Field("col").build();
+
+    PCollection<Row> inputRows =
+        pipeline
+            .apply(
+                Create.of(
+                    TestUtils.rowsBuilderOf(schema)
+                        .addRows(
+                            0, 1,
+                            1, 3,
+                            2, 4)
+                        .getRows()))
+            .setRowSchema(schema);
+
+    String sql = "SELECT key, SINGLE_VALUE(col) as single_value FROM PCOLLECTION GROUP BY key";
+
+    PCollection<Row> result = inputRows.apply("sql", SqlTransform.query(sql));
+
+    Map<Integer, List<Integer>> allowedTuples = new HashMap<>();
+    allowedTuples.put(0, Arrays.asList(1));
+    allowedTuples.put(1, Arrays.asList(3));
+    allowedTuples.put(2, Arrays.asList(4));
+
+    PAssert.that(result)
+        .satisfies(
+            input -> {
+              Iterator<Row> iter = input.iterator();
+              while (iter.hasNext()) {
+                Row row = iter.next();
+                List<Integer> values = allowedTuples.remove(row.getInt32("key"));
+                assertTrue(values != null);
+                assertTrue(values.contains(row.getInt32("single_value")));
+              }
+              assertTrue(allowedTuples.isEmpty());
+              return null;
+            });
+
+    pipeline.run();
+  }
+
+  /** GROUP-BY with the SINGLE_VALUE aggregation function fails on multiple values. */
+  @Test
+  public void testSingleValueFunction_throwsOnMultiple() throws Exception {
+    pipeline.enableAbandonedNodeEnforcement(false);
+
+    Schema schema = Schema.builder().addInt32Field("key").addInt32Field("col").build();
+
+    PCollection<Row> inputRows =
+        pipeline
+            .apply(
+                Create.of(
+                    TestUtils.rowsBuilderOf(schema)
+                        .addRows(
+                            0, 1,
+                            0, 2)
+                        .getRows()))
+            .setRowSchema(schema);
+
+    String sql = "SELECT key, SINGLE_VALUE(col) as single_value FROM PCOLLECTION GROUP BY key";
+
+    inputRows.apply("sql", SqlTransform.query(sql));
+
+    thrown.expect(Exception.class);
+    thrown.expectMessage("Subquery returned more than one row");
+    pipeline.run();
+  }
+
   @Test
   public void testBitOrFunction() throws Exception {
     pipeline.enableAbandonedNodeEnforcement(false);