Skip to content

Commit 8669518

Browse files
committed
refactor extensionarrays to use from_sequence, use df from_mgr to resolve warnings
1 parent dfcb311 commit 8669518

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

connectorx-python/connectorx/__init__.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -459,23 +459,27 @@ def reconstruct_pandas(df_infos: _DataframeInfos) -> pd.DataFrame:
459459
pd.core.internals.make_block(block_data, placement=binfo.cids)
460460
)
461461
elif binfo.dt == 1: # IntegerArray
462+
integer_array = pd.core.arrays.IntegerArray._from_sequence(block_data[0])
463+
integer_array._mask = block_data[1]
462464
blocks.append(
463465
pd.core.internals.make_block(
464-
pd.core.arrays.IntegerArray(block_data[0], block_data[1]),
466+
integer_array,
465467
placement=binfo.cids[0],
466468
)
467469
)
468470
elif binfo.dt == 2: # BooleanArray
471+
bool_array = pd.core.arrays.BooleanArray._from_sequence(block_data[0])
472+
bool_array._mask = block_data[1]
469473
blocks.append(
470474
pd.core.internals.make_block(
471-
pd.core.arrays.BooleanArray(block_data[0], block_data[1]),
475+
bool_array,
472476
placement=binfo.cids[0],
473477
)
474478
)
475479
elif binfo.dt == 3: # DatetimeArray
476480
blocks.append(
477481
pd.core.internals.make_block(
478-
pd.core.arrays.DatetimeArray(block_data), placement=binfo.cids
482+
pd.core.arrays.DatetimeArray._from_sequence(block_data), placement=binfo.cids
479483
)
480484
)
481485
else:
@@ -484,7 +488,7 @@ def reconstruct_pandas(df_infos: _DataframeInfos) -> pd.DataFrame:
484488
block_manager = pd.core.internals.BlockManager(
485489
blocks, [pd.Index(headers), pd.RangeIndex(start=0, stop=nrows, step=1)]
486490
)
487-
df = pd.DataFrame(block_manager)
491+
df = pd.DataFrame._from_mgr(block_manager, axes=[headers, range(nrows)])
488492
return df
489493

490494

0 commit comments

Comments
 (0)