Skip to content

Commit

Permalink
PERF: use new to_records() argument in to_stata()
Browse files Browse the repository at this point in the history
  • Loading branch information
qwhelan committed Jan 30, 2019
1 parent 4cbee17 commit 85df94f
Showing 1 changed file with 5 additions and 12 deletions.
17 changes: 5 additions & 12 deletions pandas/io/stata.py
Original file line number Diff line number Diff line change
Expand Up @@ -2385,8 +2385,7 @@ def _prepare_data(self):
data = self._convert_strls(data)

# 3. Convert bad string data to '' and pad to correct length
dtypes = []
data_cols = []
dtypes = {}
has_strings = False
native_byteorder = self._byteorder == _set_endianness(sys.byteorder)
for i, col in enumerate(data):
Expand All @@ -2395,22 +2394,16 @@ def _prepare_data(self):
has_strings = True
data[col] = data[col].fillna('').apply(_pad_bytes, args=(typ,))
stype = 'S{type}'.format(type=typ)
dtypes.append(('c' + str(i), stype))
string = data[col].str.encode(self._encoding)
data_cols.append(string.values.astype(stype))
dtypes[col] = stype
data[col] = data[col].str.encode(self._encoding).astype(stype)
else:
values = data[col].values
dtype = data[col].dtype
if not native_byteorder:
dtype = dtype.newbyteorder(self._byteorder)
dtypes.append(('c' + str(i), dtype))
data_cols.append(values)
dtypes = np.dtype(dtypes)
dtypes[col] = dtype

if has_strings or not native_byteorder:
self.data = np.fromiter(zip(*data_cols), dtype=dtypes)
else:
self.data = data.to_records(index=False)
self.data = data.to_records(index=False, column_dtypes=dtypes)

def _write_data(self):
data = self.data
Expand Down

0 comments on commit 85df94f

Please sign in to comment.