Skip to content

WIP/ENH: allow categoricals in msgpack #12191

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.

Sign up for GitHub

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

Closed
wants to merge 1 commit into from
Closed

WIP/ENH: allow categoricals in msgpack #12191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
29 changes: 26 additions & 3 deletions pandas/io/packers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,11 +47,12 @@
from pandas.compat import u
from pandas import (Timestamp, Period, Series, DataFrame, # noqa
Index, MultiIndex, Float64Index, Int64Index,
Panel, RangeIndex, PeriodIndex, DatetimeIndex)
Panel, RangeIndex, PeriodIndex, DatetimeIndex,
Categorical)
from pandas.sparse.api import SparseSeries, SparseDataFrame, SparsePanel
from pandas.sparse.array import BlockIndex, IntIndex
from pandas.core.generic import NDFrame
from pandas.core.common import needs_i8_conversion
from pandas.core.common import needs_i8_conversion, is_categorical_dtype
from pandas.io.common import get_filepath_or_buffer
from pandas.core.internals import BlockManager, make_block
import pandas.core.internals as internals
Expand Down Expand Up @@ -170,6 +171,7 @@ def read(fh):
# this is platform int, which we need to remap to np.int64
# for compat on windows platforms
7: np.dtype('int64'),
'category': 'category'
}


Expand Down Expand Up @@ -209,6 +211,14 @@ def convert(values):
if dtype == np.object_:
return v.tolist()

if is_categorical_dtype(dtype):
return {
'codes': {'dtype': values.codes.dtype.name,
'data': convert(values.codes)},
'categories': {'dtype': values.categories.dtype.name,
'data': convert(values.categories.values)}
}

if compressor == 'zlib':

# return string arrays like they are
Expand Down Expand Up @@ -242,6 +252,15 @@ def unconvert(values, dtype, compress=None):
if as_is_ext:
values = values.data

if is_categorical_dtype(dtype):
return Categorical.from_codes(
unconvert(values['codes']['data'],
dtype_for(values['codes']['dtype']),
compress=compress),
unconvert(values['categories']['data'],
dtype_for(values['categories']['dtype']),
compress=compress))

if dtype == np.object_:
return np.array(values, dtype=object)

Expand Down Expand Up @@ -495,11 +514,15 @@ def decode(obj):

elif typ == 'series':
dtype = dtype_for(obj['dtype'])
ctor_dtype = dtype
if is_categorical_dtype(dtype):
# Series ctor doesn't take dtype with categorical
ctor_dtype = None
index = obj['index']
return globals()[obj['klass']](unconvert(obj['data'], dtype,
obj['compress']),
index=index,
dtype=dtype,
dtype=ctor_dtype,
name=obj['name'])
elif typ == 'block_manager':
axes = obj['axes']
Expand Down
7 changes: 5 additions & 2 deletions pandas/io/tests/test_packers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from pandas import compat
from pandas.compat import u
from pandas import (Series, DataFrame, Panel, MultiIndex, bdate_range,
date_range, period_range, Index)
date_range, period_range, Index, Categorical)
from pandas.io.packers import to_msgpack, read_msgpack
import pandas.util.testing as tm
from pandas.util.testing import (ensure_clean, assert_index_equal,
Expand Down Expand Up @@ -330,11 +330,13 @@ def setUp(self):
'C': ['foo1', 'foo2', 'foo3', 'foo4', 'foo5'],
'D': date_range('1/1/2009', periods=5),
'E': [0., 1, Timestamp('20100101'), 'foo', 2.],
'F': Categorical(['a', 'b', 'c', 'd', 'e'])
}

self.d['float'] = Series(data['A'])
self.d['int'] = Series(data['B'])
self.d['mixed'] = Series(data['E'])
self.d['categorical'] = Series(data['F'])

def test_basic(self):

Expand All @@ -356,13 +358,14 @@ def setUp(self):
'C': ['foo1', 'foo2', 'foo3', 'foo4', 'foo5'],
'D': date_range('1/1/2009', periods=5),
'E': [0., 1, Timestamp('20100101'), 'foo', 2.],
'F': Categorical(['a', 'b', 'c', 'd', 'e'])
}

self.frame = {
'float': DataFrame(dict(A=data['A'], B=Series(data['A']) + 1)),
'int': DataFrame(dict(A=data['B'], B=Series(data['B']) + 1)),
'mixed': DataFrame(dict([(k, data[k])
for k in ['A', 'B', 'C', 'D']]))}
for k in ['A', 'B', 'C', 'D', 'F']]))}

self.panel = {
'float': Panel(dict(ItemA=self.frame['float'],
Expand Down