@@ -1426,28 +1426,21 @@ def _groupby_op(
1426
1426
return self ._maybe_mask_result (res_values , result_mask )
1427
1427
1428
1428
1429
- def transpose_homogenous_masked_arrays (masked_arrays ):
1429
+ def transpose_homogenous_masked_arrays (masked_arrays ) -> list [ BaseMaskedArray ] :
1430
1430
"""Transpose masked arrays in a list, but faster.
1431
1431
1432
1432
Input should be a list of 1-dim masked arrays of equal length and all have the
1433
1433
same dtype. The caller is responsible for ensuring validity of input data.
1434
1434
"""
1435
- transposed_shape = len (masked_arrays ), len (masked_arrays [0 ])
1436
- values = [arr ._data for arr in masked_arrays ]
1437
- transposed_values = np .empty (transposed_shape , dtype = values [0 ].dtype )
1438
- for i , val in enumerate (values ):
1439
- transposed_values [i , :] = val
1440
- transposed_values = transposed_values .copy (order = "F" )
1441
-
1442
- masks = [arr ._mask for arr in masked_arrays ]
1443
- transposed_masks = np .empty (transposed_shape , dtype = masks [0 ].dtype )
1444
- for i , mask in enumerate (masks ):
1445
- transposed_masks [i , :] = mask
1446
- transposed_masks = transposed_masks .copy (order = "F" )
1435
+ values = [arr ._data .reshape (1 , - 1 ) for arr in masked_arrays ]
1436
+ transposed_values = np .concatenate (values , axis = 0 )
1437
+
1438
+ masks = [arr ._mask .reshape (1 , - 1 ) for arr in masked_arrays ]
1439
+ transposed_masks = np .concatenate (masks , axis = 0 )
1447
1440
1448
1441
dtype = masked_arrays [0 ].dtype
1449
1442
arr_type = dtype .construct_array_type ()
1450
- transposed_arrays = []
1443
+ transposed_arrays : list [ BaseMaskedArray ] = []
1451
1444
for i in range (transposed_values .shape [1 ]):
1452
1445
transposed_arr = arr_type (transposed_values [:, i ], mask = transposed_masks [:, i ])
1453
1446
transposed_arrays .append (transposed_arr )
0 commit comments