puppoc

Author: polyzos

.idea/icon.png

@@ -73,6 +73,9 @@ public class FlussSinkBuilder<InputT> {
     private final Map<String, String> configOptions = new HashMap<>();
     private FlussSerializationSchema<InputT> serializationSchema;
     private boolean shuffleByBucketId = true;
+    // Optional list of columns for partial update. When set, upsert will only update these columns.
+    // The primary key columns must be fully specified in this list.
+    private List<String> partialUpdateColumns;
 
     /** Set the bootstrap server for the sink. */
     public FlussSinkBuilder<InputT> setBootstrapServers(String bootstrapServers) {
@@ -98,6 +101,28 @@ public FlussSinkBuilder<InputT> setShuffleByBucketId(boolean shuffleByBucketId)
         return this;
     }
 
+    /**
+     * Enable partial update by specifying the column names to update for upsert tables.
+     * Primary key columns must be included in this list.
+     */
+    public FlussSinkBuilder<InputT> setPartialUpdateColumns(List<String> columns) {
+        this.partialUpdateColumns = columns;
+        return this;
+    }
+
+    /**
+     * Enable partial update by specifying the column names to update for upsert tables.
+     * Convenience varargs overload.
+     */
+    public FlussSinkBuilder<InputT> setPartialUpdateColumns(String... columns) {
+        if (columns == null) {
+            this.partialUpdateColumns = null;
+        } else {
+            this.partialUpdateColumns = java.util.Arrays.asList(columns);
+        }
+        return this;
+    }
+
     /** Set a configuration option. */
     public FlussSinkBuilder<InputT> setOption(String key, String value) {
         configOptions.put(key, value);
@@ -151,14 +176,24 @@ public FlussSink<InputT> build() {
 
         boolean isUpsert = tableInfo.hasPrimaryKey();
 
+        int[] targetColumnIndexes = null;
+        if (isUpsert) {
+            // Compute target column indexes for partial update if columns were provided
+            targetColumnIndexes =
+                    computeTargetColumnIndexes(
+                            tableRowType.getFieldNames(),
+                            tableInfo.getPrimaryKeys(),
+                            partialUpdateColumns);
+        }
+
         if (isUpsert) {
             LOG.info("Initializing Fluss upsert sink writer ...");
             writerBuilder =
                     new FlinkSink.UpsertSinkWriterBuilder<>(
                             tablePath,
                             flussConfig,
                             tableRowType,
-                            null, // not support partialUpdateColumns yet
+                            targetColumnIndexes,
                             numBucket,
                             bucketKeys,
                             partitionKeys,
@@ -193,4 +228,47 @@ private void validateConfiguration() {
         checkNotNull(tableName, "Table name is required but not provided.");
         checkArgument(!tableName.isEmpty(), "Table name cannot be empty.");
     }
+
+    // -------------- Test-visible helper methods --------------
+    /**
+     * Computes target column indexes for partial updates.
+     * If {@code specifiedColumns} is null or empty, returns null indicating full update.
+     * Validates that all primary key columns are included in the specified columns.
+     *
+     * @param allFieldNames the list of all field names in table row type order
+     * @param primaryKeyNames the list of primary key column names
+     * @param specifiedColumns the optional list of columns specified for partial update
+     * @return the indexes into {@code allFieldNames} corresponding to {@code specifiedColumns}, or null for full update
+     * @throws IllegalArgumentException if a specified column does not exist or primary key coverage is incomplete
+     */
+    static int[] computeTargetColumnIndexes(
+            List<String> allFieldNames, List<String> primaryKeyNames, List<String> specifiedColumns) {
+        if (specifiedColumns == null || specifiedColumns.isEmpty()) {
+            return null; // full update
+        }
+
+        // Map specified column names to indexes
+        int[] indexes = new int[specifiedColumns.size()];
+        for (int i = 0; i < specifiedColumns.size(); i++) {
+            String col = specifiedColumns.get(i);
+            int idx = allFieldNames.indexOf(col);
+            checkArgument(
+                    idx >= 0,
+                    "Column '%s' not found in table schema: %s",
+                    col,
+                    allFieldNames);
+            indexes[i] = idx;
+        }
+
+        // Validate that all primary key columns are covered
+        for (String pk : primaryKeyNames) {
+            checkArgument(
+                    specifiedColumns.contains(pk),
+                    "Partial updates must include all primary key columns. Missing primary key column: %s. Provided columns: %s",
+                    pk,
+                    specifiedColumns);
+        }
+
+        return indexes;
+    }
 }

.idea/vcs.xml

@@ -24,6 +24,7 @@
 import org.junit.jupiter.api.Test;
 
 import java.lang.reflect.Field;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -171,12 +172,57 @@ void testFluentChaining() {
                         .setTable(tableName)
                         .setOption("key1", "value1")
                         .setOptions(new HashMap<>())
-                        .setShuffleByBucketId(false);
+                        .setShuffleByBucketId(false)
+                        .setPartialUpdateColumns("id", "price");
 
         // Verify the builder instance is returned
         assertThat(chainedBuilder).isInstanceOf(FlussSinkBuilder.class);
     }
 
+    @Test
+    void testComputeTargetColumnIndexesFullUpdate() {
+        int[] result =
+                FlussSinkBuilder.computeTargetColumnIndexes(
+                        Arrays.asList("id", "name", "price"),
+                        Arrays.asList("id"),
+                        null);
+        assertThat(result).isNull();
+    }
+
+    @Test
+    void testComputeTargetColumnIndexesValidPartialIncludesPk() {
+        int[] result =
+                FlussSinkBuilder.computeTargetColumnIndexes(
+                        Arrays.asList("id", "name", "price", "ts"),
+                        Arrays.asList("id"),
+                        Arrays.asList("id", "price"));
+        assertThat(result).containsExactly(0, 2);
+    }
+
+    @Test
+    void testComputeTargetColumnIndexesMissingPkThrows() {
+        assertThatThrownBy(
+                        () ->
+                                FlussSinkBuilder.computeTargetColumnIndexes(
+                                        Arrays.asList("id", "name", "price"),
+                                        Arrays.asList("id"),
+                                        Arrays.asList("name", "price")))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessageContaining("Partial updates must include all primary key columns");
+    }
+
+    @Test
+    void testComputeTargetColumnIndexesUnknownColumnThrows() {
+        assertThatThrownBy(
+                        () ->
+                                FlussSinkBuilder.computeTargetColumnIndexes(
+                                        Arrays.asList("id", "name"),
+                                        Arrays.asList("id"),
+                                        Arrays.asList("id", "unknown")))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessageContaining("not found in table schema");
+    }
+
     // Helper method to get private field values using reflection
     @SuppressWarnings("unchecked")
     private <T> T getFieldValue(Object object, String fieldName) throws Exception {