annotation fix

Author: rok

python/pyarrow-stubs/pyarrow/compute.pyi

@@ -654,7 +654,7 @@ def round(
 def round_to_multiple(
     x: _NumericScalarT | _NumericArrayT | list | Expression,
     /,
-    multiple: int = 0,
+    multiple: int | float | NumericScalar = 1.0,
     round_mode: Literal[
         "down",
         "up",

python/pyarrow/tests/parquet/test_dataset.py

@@ -971,7 +971,7 @@ def _test_write_to_dataset_with_partitions(base_path,
     input_df_cols = input_df.columns.tolist()
     assert partition_by == input_df_cols[-1 * len(partition_by):]
 
-    input_df = input_df[cols]  # type: ignore[assignment]
+    input_df = input_df.loc[:, cols]
     # Partitioned columns become 'categorical' dtypes
     for col in partition_by:
         output_df[col] = output_df[col].astype('category')