Skip to content

Commit

Permalink
ENH: Add string_as_bytes option for df.to_records() (pandas-dev#18146)
Browse files Browse the repository at this point in the history
This options records dtype for string as arrays as 'Sx', where x
is the length of the longest string, instead of 'O"
  • Loading branch information
qinghao1 authored and Chu Qinghao committed Aug 7, 2018
1 parent 2156431 commit 2f48a8f
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 2 deletions.
34 changes: 32 additions & 2 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -1335,7 +1335,7 @@ def from_records(cls, data, index=None, exclude=None, columns=None,

return cls(mgr)

def to_records(self, index=True, convert_datetime64=None):
def to_records(self, index=True, convert_datetime64=None, strings_as_bytes=False):
"""
Convert DataFrame to a NumPy record array.
Expand All @@ -1351,6 +1351,9 @@ def to_records(self, index=True, convert_datetime64=None):
Whether to convert the index to datetime.datetime if it is a
DatetimeIndex.
strings_as_bytes: boolean, default False
Store strings as bytes ('S' dtype) instead of Python objects
('O' dtype)
Returns
-------
Expand Down Expand Up @@ -1401,6 +1404,24 @@ def to_records(self, index=True, convert_datetime64=None):
rec.array([('2018-01-01T09:00:00.000000000', 1, 0.5 ),
('2018-01-01T09:01:00.000000000', 2, 0.75)],
dtype=[('index', '<M8[ns]'), ('A', '<i8'), ('B', '<f8')])
By default, strings are recorded as dtype `O` for object:
>>> df = pd.DataFrame({'A': [1, 2], 'B': ['abc', 'defg']},
... index=['a', 'b'])
>>> df.to_records()
rec.array([('a', 1, 'abc'), ('b', 2, 'defg')],
dtype=[('index', 'O'), ('A', '<i8'), ('B', 'O')])
This can be inefficient (e.g. for short strings, or when storing with
`np.save()`). They can be recorded as dtype `S` for zero-terminated
bytes instead:
>>> df = pd.DataFrame({'A': [1, 2], 'B': ['abc', 'defg']},
... index=['a', 'b'])
>>> df.to_records()
rec.array([('a', 1, 'abc'), ('b', 2, 'defg')],
dtype=[('index', 'O'), ('A', '<i8'), ('B', 'S4')])
"""

if convert_datetime64 is not None:
Expand Down Expand Up @@ -1436,7 +1457,16 @@ def to_records(self, index=True, convert_datetime64=None):
arrays = [self[c].get_values() for c in self.columns]
names = lmap(compat.text_type, self.columns)

formats = [v.dtype for v in arrays]
if strings_as_bytes:
# GH18146
# for string arrays, set dtype as zero-terminated bytes with max
# length equals to that of the longest string
formats = ['S{}'.format(max(lmap(len, v)))
if v.dtype == '|O'
else v.dtype
for v in arrays]
else:
formats = [v.dtype for v in arrays]
return np.rec.fromarrays(
arrays,
dtype={'names': names, 'formats': formats}
Expand Down
11 changes: 11 additions & 0 deletions pandas/tests/frame/test_convert_to.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,17 @@ def test_to_records_with_categorical(self):
dtype=[('index', '=i8'), ('0', 'O')])
tm.assert_almost_equal(result, expected)

def test_to_records_with_strings_as_bytes(self):

# GH18146

df = DataFrame({'A': [1, 2], 'B': ['abc', 'defg']},
index=['a', 'b'])
result = df.to_records(strings_as_bytes=True)
expected = np.rec.array([('a', 1, 'abc'), ('b', 2, 'defg')],
dtype=[('index', 'O'), ('A', '<i8'), ('B', 'S4')])
tm.assert_almost_equal(result, expected)

@pytest.mark.parametrize('mapping', [
dict,
collections.defaultdict(list),
Expand Down

0 comments on commit 2f48a8f

Please sign in to comment.