Skip to content

Commit dd3771b

Browse files
committed
PERF: use new to_records() argument in to_stata()
1 parent 4cbee17 commit dd3771b

File tree

1 file changed

+5
-15
lines changed

1 file changed

+5
-15
lines changed

pandas/io/stata.py

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2385,32 +2385,22 @@ def _prepare_data(self):
23852385
data = self._convert_strls(data)
23862386

23872387
# 3. Convert bad string data to '' and pad to correct length
2388-
dtypes = []
2389-
data_cols = []
2390-
has_strings = False
2388+
dtypes = {}
23912389
native_byteorder = self._byteorder == _set_endianness(sys.byteorder)
23922390
for i, col in enumerate(data):
23932391
typ = typlist[i]
23942392
if typ <= self._max_string_length:
2395-
has_strings = True
23962393
data[col] = data[col].fillna('').apply(_pad_bytes, args=(typ,))
23972394
stype = 'S{type}'.format(type=typ)
2398-
dtypes.append(('c' + str(i), stype))
2399-
string = data[col].str.encode(self._encoding)
2400-
data_cols.append(string.values.astype(stype))
2395+
dtypes[col] = stype
2396+
data[col] = data[col].str.encode(self._encoding).astype(stype)
24012397
else:
2402-
values = data[col].values
24032398
dtype = data[col].dtype
24042399
if not native_byteorder:
24052400
dtype = dtype.newbyteorder(self._byteorder)
2406-
dtypes.append(('c' + str(i), dtype))
2407-
data_cols.append(values)
2408-
dtypes = np.dtype(dtypes)
2401+
dtypes[col] = dtype
24092402

2410-
if has_strings or not native_byteorder:
2411-
self.data = np.fromiter(zip(*data_cols), dtype=dtypes)
2412-
else:
2413-
self.data = data.to_records(index=False)
2403+
self.data = data.to_records(index=False, column_dtypes=dtypes)
24142404

24152405
def _write_data(self):
24162406
data = self.data

0 commit comments

Comments
 (0)