Skip to content

Implement BlockManager.iset #32350

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Mar 3, 2020
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 24 additions & 10 deletions pandas/core/internals/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import numpy as np

from pandas._libs import Timedelta, Timestamp, internals as libinternals, lib
from pandas._typing import DtypeObj
from pandas._typing import DtypeObj, Label
from pandas.util._validators import validate_bool_kwarg

from pandas.core.dtypes.cast import (
Expand Down Expand Up @@ -996,7 +996,25 @@ def delete(self, item):
)
self._rebuild_blknos_and_blklocs()

def set(self, item, value):
def set(self, item: Label, value):
"""
Set new item in-place.

Notes
-----
Does not consolidate.
Adds new Block if not contained in the current items Index.
"""
try:
loc = self.items.get_loc(item)
except KeyError:
# This item wasn't present, just insert at end
self.insert(len(self.items), item, value)
return

self.iset(loc, value)

def iset(self, loc: Union[int, slice, np.ndarray], value):
"""
Set new item in-place. Does not consolidate. Adds new Block if not
contained in the current set of items
Expand Down Expand Up @@ -1029,13 +1047,6 @@ def value_getitem(placement):
"Shape of new values must be compatible with manager shape"
)

try:
loc = self.items.get_loc(item)
except KeyError:
# This item wasn't present, just insert at end
self.insert(len(self.items), item, value)
return

if isinstance(loc, int):
loc = [loc]

Expand Down Expand Up @@ -1081,7 +1092,7 @@ def value_getitem(placement):
unfit_mgr_locs = np.concatenate(unfit_mgr_locs)
unfit_count = len(unfit_mgr_locs)

new_blocks = []
new_blocks: List[Block] = []
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rather than add the variable annotation, does new_blocks need to be initialized outside the if else? make_block has the correct return type and the extend wouldn't be needed. (but I guess this mypy fixup is orthogonal to this PR)

Copy link
Member

@simonjayhawkins simonjayhawkins Mar 3, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not a fan of variable annotations when not required. but not a blocker.

if value_is_extension_type:
# This code (ab-)uses the fact that sparse blocks contain only
# one item.
Expand Down Expand Up @@ -1140,6 +1151,9 @@ def insert(self, loc: int, item, value, allow_duplicates: bool = False):
# insert to the axis; this could possibly raise a TypeError
new_axis = self.items.insert(loc, item)

if value.ndim == self.ndim - 1 and not is_extension_array_dtype(value):
value = _safe_reshape(value, (1,) + value.shape)

block = make_block(values=value, ndim=self.ndim, placement=slice(loc, loc + 1))

for blkno, count in _fast_count_smallints(self._blknos[loc:]):
Expand Down