@@ -305,8 +305,11 @@ def arr2shm(self, vals):
305
305
"""Array to shared memory. Returns (shm_name, shape, dtype) used for restore."""
306
306
assert vals .ndim == 1 , (vals .ndim , vals .shape , vals )
307
307
shm = self .SharedMemory (size = vals .nbytes , create = True )
308
- buf = np .ndarray (vals .shape , dtype = vals .dtype , buffer = shm .buf )
309
- buf [:] = vals [:] # Copy into shared memory
308
+ # np.array can't handle pandas' tz-aware datetimes
309
+ # https://github.com/numpy/numpy/issues/18279
310
+ buf = np .ndarray (vals .shape , dtype = vals .dtype .base , buffer = shm .buf )
311
+ has_tz = getattr (vals .dtype , 'tz' , None )
312
+ buf [:] = vals .tz_localize (None ) if has_tz else vals # Copy into shared memory
310
313
return shm .name , vals .shape , vals .dtype
311
314
312
315
def df2shm (self , df ):
@@ -316,18 +319,18 @@ def df2shm(self, df):
316
319
))
317
320
318
321
@staticmethod
319
- def shm2arr (shm , shape , dtype ):
320
- arr = np .ndarray (shape , dtype = dtype , buffer = shm .buf )
322
+ def shm2s (shm , shape , dtype ) -> pd . Series :
323
+ arr = np .ndarray (shape , dtype = dtype . base , buffer = shm .buf )
321
324
arr .setflags (write = False )
322
- return arr
325
+ return pd . Series ( arr , dtype = dtype )
323
326
324
327
_DF_INDEX_COL = '__bt_index'
325
328
326
329
@staticmethod
327
330
def shm2df (data_shm ):
328
331
shm = [SharedMemory (name = name , create = False , track = False ) for _ , name , _ , _ in data_shm ]
329
332
df = pd .DataFrame ({
330
- col : SharedMemoryManager .shm2arr (shm , shape , dtype )
333
+ col : SharedMemoryManager .shm2s (shm , shape , dtype )
331
334
for shm , (col , _ , shape , dtype ) in zip (shm , data_shm )})
332
335
df .set_index (SharedMemoryManager ._DF_INDEX_COL , drop = True , inplace = True )
333
336
df .index .name = None
0 commit comments