[rust][client][server] Require only omitted columns to be nullable in partial update

fluss-client/src/main/java/org/apache/fluss/client/table/writer/UpsertWrite.java

@@ -54,8 +54,10 @@ private UpsertWrite(@Nullable int[] targetColumns) {
      * row will be removed when all columns except primary key are null after a {@link
      * UpsertWriter#delete(InternalRow)} operation.
      *
-     * <p>Note: The specified columns must be a contains all columns of primary key, and all columns
-     * except primary key should be nullable.
+     * <p>Note: The specified columns must contain all columns of the primary key, and all columns
+     * omitted from the specified columns should be nullable, since they are set to null when the
+     * row doesn't exist. A {@link UpsertWriter#delete(InternalRow)} operation additionally requires
+     * the specified columns except primary key to be nullable, since it sets them to null.
      *
      * @param targetColumns the columns to partial update,
      */

fluss-client/src/main/java/org/apache/fluss/client/table/writer/UpsertWriterImpl.java

@@ -126,7 +126,6 @@ private static void sanityCheck(
             targetColumnsSet.set(targetColumnIndex);
         }
 
-        BitSet pkColumnSet = new BitSet();
         // check the target columns contains the primary key
         for (String key : primaryKeys) {
             int pkIndex = rowType.getFieldIndex(key);
@@ -136,7 +135,6 @@ private static void sanityCheck(
                                 "The target write columns %s must contain the primary key columns %s.",
                                 rowType.project(targetColumns).getFieldNames(), primaryKeys));
             }
-            pkColumnSet.set(pkIndex);
         }
 
         BitSet autoIncrementColumnSet = new BitSet();
@@ -152,17 +150,21 @@ private static void sanityCheck(
             autoIncrementColumnSet.set(autoIncrementColumnIndex);
         }
 
-        // check the columns not in targetColumns should be nullable
+        // an omitted column is written as null, so it must be nullable. auto increment columns
+        // are always omitted and only get their value on the server.
         for (int i = 0; i < rowType.getFieldCount(); i++) {
-            // column not in primary key and not in auto increment column
-            if (!pkColumnSet.get(i) && !autoIncrementColumnSet.get(i)) {
-                // the column should be nullable
-                if (!rowType.getTypeAt(i).isNullable()) {
+            if (!targetColumnsSet.get(i) && !rowType.getTypeAt(i).isNullable()) {
+                String columnName = rowType.getFieldNames().get(i);
+                if (autoIncrementColumnSet.get(i)) {
                     throw new IllegalArgumentException(
                             String.format(
-                                    "Partial Update requires all columns except primary key to be nullable, but column %s is NOT NULL.",
-                                    rowType.getFieldNames().get(i)));
+                                    "Partial Update requires the auto increment column %s to be nullable, since it is always omitted from the target columns and assigned by the server.",
+                                    columnName));
                 }
+                throw new IllegalArgumentException(
+                        String.format(
+                                "Partial Update requires all columns omitted from the target columns to be nullable, but omitted column %s is NOT NULL.",
+                                columnName));
             }
         }
     }

fluss-client/src/test/java/org/apache/fluss/client/table/FlussTableITCase.java

@@ -34,6 +34,7 @@
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.Configuration;
 import org.apache.fluss.config.MemorySize;
+import org.apache.fluss.exception.InvalidTargetColumnException;
 import org.apache.fluss.fs.FsPath;
 import org.apache.fluss.fs.TestFileSystem;
 import org.apache.fluss.metadata.DataLakeFormat;
@@ -90,6 +91,7 @@
 import static org.apache.fluss.testutils.DataTestUtils.row;
 import static org.apache.fluss.testutils.InternalRowAssert.assertThatRow;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** IT case for {@link FlussTable}. */
@@ -955,6 +957,45 @@ void testPartialPutAndDelete() throws Exception {
         table.close();
     }
 
+    @Test
+    void testPartialPutWithNotNullTargetColumn() throws Exception {
+        Schema schema =
+                Schema.newBuilder()
+                        .column("a", DataTypes.INT())
+                        .column("b", DataTypes.STRING())
+                        .column("c", new BigIntType(false))
+                        .primaryKey("a")
+                        .build();
+        TableDescriptor tableDescriptor =
+                TableDescriptor.builder().schema(schema).distributedBy(3, "a").build();
+        createTable(DATA1_TABLE_PATH_PK, tableDescriptor, true);
+
+        try (Table table = conn.getTable(DATA1_TABLE_PATH_PK)) {
+            // seed a full row so that column b, outside the target columns, is not null
+            table.newUpsert().createWriter().upsert(row(1, "bbb", 1L)).get();
+
+            UpsertWriter upsertWriter = table.newUpsert().partialUpdate("a", "c").createWriter();
+            upsertWriter.upsert(row(1, null, 100L)).get();
+
+            Lookuper lookuper = table.newLookup().createLookuper();
+            GenericRow rowKey = row(1);
+            assertThat(lookupRow(lookuper, rowKey))
+                    .isEqualTo(compactedRow(schema.getRowType(), new Object[] {1, "bbb", 100L}));
+
+            // updating the same target columns again keeps column b untouched
+            upsertWriter.upsert(row(1, null, 200L)).get();
+            assertThat(lookupRow(lookuper, rowKey))
+                    .isEqualTo(compactedRow(schema.getRowType(), new Object[] {1, "bbb", 200L}));
+
+            // a partial delete would null out the NOT NULL column c, and only the server can
+            // tell, since it depends on the stored row
+            assertThatThrownBy(() -> upsertWriter.delete(row(1, null, 200L)).get())
+                    .cause()
+                    .isInstanceOf(InvalidTargetColumnException.class)
+                    .hasMessageContaining("but target column c is NOT NULL.");
+        }
+    }
+
     @Test
     void testInvalidPartialUpdate() throws Exception {
         Schema schema =
@@ -975,13 +1016,15 @@ void testInvalidPartialUpdate() throws Exception {
                     .hasMessage(
                             "The target write columns [b] must contain the primary key columns [a].");
 
-            // the column not in the primary key is nullable, should throw exception
+            // the column omitted from the target columns is NOT NULL, should throw exception
             assertThatThrownBy(() -> table.newUpsert().partialUpdate("a", "b").createWriter())
                     .hasMessage(
-                            "Partial Update requires all columns except primary key to be nullable, but column c is NOT NULL.");
-            assertThatThrownBy(() -> table.newUpsert().partialUpdate("a", "c").createWriter())
-                    .hasMessage(
-                            "Partial Update requires all columns except primary key to be nullable, but column c is NOT NULL.");
+                            "Partial Update requires all columns omitted from the target columns to be nullable, "
+                                    + "but omitted column c is NOT NULL.");
+            // column c is NOT NULL but listed in the target columns, so it is always provided by
+            // the writer and should be accepted
+            assertThatCode(() -> table.newUpsert().partialUpdate("a", "c").createWriter())
+                    .doesNotThrowAnyException();
             assertThatThrownBy(() -> table.newUpsert().partialUpdate("a", "d").createWriter())
                     .hasMessage(
                             "Can not find target column: d for table test_db_1.test_pk_table_1.");

fluss-rust/crates/fluss/src/client/table/upsert.rs

@@ -225,8 +225,6 @@ impl UpsertWriterFactory {
             target_column_set.set(target_index, true);
         }
 
-        let mut pk_column_set = bitvec![0; field_count];
-
         // check the target columns contains the primary key
         for primary_key in primary_keys {
             let pk_index = row_type.get_field_index(primary_key.as_str());
@@ -241,7 +239,6 @@ impl UpsertWriterFactory {
                             ),
                         });
                     }
-                    pk_column_set.set(pk_index, true);
                 }
                 None => {
                     return Err(IllegalArgument {
@@ -267,24 +264,28 @@ impl UpsertWriterFactory {
                         ),
                     });
                 }
-
                 auto_increment_column_set.set(index, true);
             }
         }
 
-        // check the columns not in targetColumns should be nullable
-        for i in 0..field_count {
-            // column not in primary key and not in auto increment column
-            if !pk_column_set[i] && !auto_increment_column_set[i] {
-                // the column should be nullable
-                if !row_type.fields().get(i).unwrap().data_type.is_nullable() {
+        // an omitted column is written as NULL, so it must be nullable. auto increment columns
+        // are always omitted and only get their value on the server.
+        for (i, field) in row_type.fields().iter().enumerate() {
+            if !target_column_set[i] && !field.data_type.is_nullable() {
+                if auto_increment_column_set[i] {
                     return Err(IllegalArgument {
                         message: format!(
-                            "Partial Update requires all columns except primary key to be nullable, but column {} is NOT NULL.",
-                            row_type.fields().get(i).unwrap().name()
+                            "Partial Update requires the auto increment column {} to be nullable, since it is always omitted from the target columns and assigned by the server.",
+                            field.name()
                         ),
                     });
                 }
+                return Err(IllegalArgument {
+                    message: format!(
+                        "Partial Update requires all columns omitted from the target columns to be nullable, but omitted column {} is NOT NULL.",
+                        field.name()
+                    ),
+                });
             }
         }
 
@@ -534,8 +535,80 @@ mod tests {
         );
 
         assert!(result.unwrap_err().to_string().contains(
-            "Partial Update requires all columns except primary key to be nullable, but column required_field is NOT NULL."
+            "Partial Update requires all columns omitted from the target columns to be nullable, but omitted column required_field is NOT NULL."
+        ));
+    }
+
+    #[test]
+    fn not_null_column_in_target_columns_is_accepted() {
+        let row_type = RowType::new(vec![
+            DataField::new("id", DataTypes::int().as_non_nullable(), None),
+            DataField::new(
+                "required_field",
+                DataTypes::string().as_non_nullable(),
+                None,
+            ),
+            DataField::new("optional_field", DataTypes::int(), None),
+        ]);
+        let primary_keys = vec!["id".to_string()];
+        let auto_increment_col_names = vec![];
+
+        // `required_field` is NOT NULL but listed as a target column, so it is always written
+        let target_columns = Some(Arc::new(vec![0usize, 1]));
+        let result = UpsertWriterFactory::sanity_check(
+            &row_type,
+            &primary_keys,
+            &auto_increment_col_names,
+            &target_columns,
+        );
+        assert!(result.is_ok());
+
+        // Listing every column is equally valid.
+        let target_columns = Some(Arc::new(vec![0usize, 1, 2]));
+        let result = UpsertWriterFactory::sanity_check(
+            &row_type,
+            &primary_keys,
+            &auto_increment_col_names,
+            &target_columns,
+        );
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn omitted_not_null_auto_increment_column_is_rejected() {
+        let row_type = RowType::new(vec![
+            DataField::new("id", DataTypes::int().as_non_nullable(), None),
+            DataField::new("name", DataTypes::string(), None),
+            DataField::new("seq", DataTypes::bigint().as_non_nullable(), None),
+        ]);
+        let primary_keys = vec!["id".to_string()];
+        let auto_increment_col_names = vec!["seq".to_string()];
+        let target_columns = Some(Arc::new(vec![0usize, 1]));
+
+        let result = UpsertWriterFactory::sanity_check(
+            &row_type,
+            &primary_keys,
+            &auto_increment_col_names,
+            &target_columns,
+        );
+
+        assert!(result.unwrap_err().to_string().contains(
+            "Partial Update requires the auto increment column seq to be nullable, since it is always omitted from the target columns and assigned by the server."
         ));
+
+        // The same table with a nullable auto increment column is fine.
+        let row_type = RowType::new(vec![
+            DataField::new("id", DataTypes::int().as_non_nullable(), None),
+            DataField::new("name", DataTypes::string(), None),
+            DataField::new("seq", DataTypes::bigint(), None),
+        ]);
+        let result = UpsertWriterFactory::sanity_check(
+            &row_type,
+            &primary_keys,
+            &auto_increment_col_names,
+            &target_columns,
+        );
+        assert!(result.is_ok());
     }
 }