Use smallest data type for partition column in P2P (#8479)

Author: fjetter

distributed/shuffle/_arrow.py

@@ -102,7 +102,11 @@ def _create_input_partition_id_array(
         table: pa.Table, input_partition_id: int
     ) -> pa.ChunkedArray:
         arrays = (
-            np.full((batch.num_rows,), input_partition_id)
+            np.full(
+                (batch.num_rows,),
+                input_partition_id,
+                dtype=np.uint32(),
+            )
             for batch in table.to_batches()
         )
         return pa.chunked_array(arrays)