[python] Validate write column schema types

paimon-python/pypaimon/tests/write/table_write_test.py

@@ -568,6 +568,47 @@ def test_dynamic_bucket_write(self):
             actual.sort_by(sort_keys),
         )
 
+    def test_column_subset_write_rejects_int64_for_int32(self):
+        pa_schema = pa.schema([
+            ('id', pa.int32()),
+            ('name', pa.string()),
+        ])
+        schema = Schema.from_pyarrow_schema(pa_schema)
+        self.catalog.create_table(
+            'default.test_column_subset_reject_int64', schema, False)
+        table = self.catalog.get_table(
+            'default.test_column_subset_reject_int64')
+
+        write_builder = table.new_batch_write_builder()
+        table_write = write_builder.new_write().with_write_type(['id'])
+        with self.assertRaises(ValueError) as e:
+            table_write.write_arrow(pa.Table.from_pydict(
+                {'id': [1]}))
+        self.assertTrue(str(e.exception).startswith(
+            "Input schema isn't consistent with table schema and write cols."))
+
+    def test_validate_schema_allows_binary_family_for_write_cols(self):
+        pa_schema = pa.schema([
+            ('id', pa.int32()),
+            ('payload', pa.binary()),
+        ])
+        schema = Schema.from_pyarrow_schema(pa_schema)
+        self.catalog.create_table(
+            'default.test_validate_binary_family', schema, False)
+        table = self.catalog.get_table('default.test_validate_binary_family')
+
+        write_builder = table.new_batch_write_builder()
+        table_write = write_builder.new_write()
+        table_write._validate_pyarrow_schema(pa.schema([
+            ('id', pa.int32()),
+            ('payload', pa.binary(4)),
+        ]))
+
+        table_write.with_write_type(['payload'])
+        table_write._validate_pyarrow_schema(pa.schema([
+            ('payload', pa.binary(4)),
+        ]))
+
     @parameterized.expand([('parquet',), ('orc',), ('avro',)])
     def test_write_time_type(self, file_format):
         time_schema = pa.schema([

paimon-python/pypaimon/write/table_write.py

@@ -154,23 +154,39 @@ def abort(self):
         self.file_store_write.abort()
 
     def _validate_pyarrow_schema(self, data_schema: pa.Schema):
-        if data_schema == self.table_pyarrow_schema:
+        if self._is_compatible_pyarrow_schema(data_schema, self.table_pyarrow_schema):
             return
-        if data_schema.names == self.file_store_write.write_cols:
-            return
-        # Allow compatible binary types: binary, fixed_size_binary[N] are interchangeable
-        if data_schema.names == self.table_pyarrow_schema.names:
-            compatible = True
-            for i in range(len(data_schema)):
-                input_type = data_schema.field(i).type
-                table_type = self.table_pyarrow_schema.field(i).type
-                if input_type != table_type:
-                    if self._is_binary_family(input_type) and self._is_binary_family(table_type):
-                        continue
-                    compatible = False
-                    break
-            if compatible:
+
+        write_cols = self.file_store_write.write_cols
+        if write_cols is not None:
+            write_cols_schema = self._write_cols_pyarrow_schema(write_cols)
+            if self._is_compatible_pyarrow_schema(data_schema, write_cols_schema):
                 return
+
+        self._raise_inconsistent_schema(data_schema)
+
+    def _is_compatible_pyarrow_schema(
+            self, data_schema: pa.Schema, expected_schema: pa.Schema) -> bool:
+        # Allow compatible binary types: binary, fixed_size_binary[N] are interchangeable
+        if data_schema.names != expected_schema.names:
+            return False
+        for i in range(len(data_schema)):
+            input_type = data_schema.field(i).type
+            expected_type = expected_schema.field(i).type
+            if input_type == expected_type:
+                continue
+            if self._is_binary_family(input_type) and self._is_binary_family(expected_type):
+                continue
+            return False
+        return True
+
+    def _write_cols_pyarrow_schema(self, write_cols: List[str]) -> pa.Schema:
+        table_fields = {
+            field.name: field for field in self.table_pyarrow_schema
+        }
+        return pa.schema([table_fields[col] for col in write_cols])
+
+    def _raise_inconsistent_schema(self, data_schema: pa.Schema):
         raise ValueError(f"Input schema isn't consistent with table schema and write cols. "
                          f"Input schema is: {data_schema} "
                          f"Table schema is: {self.table_pyarrow_schema} "