Skip to content

Commit bf8f329

Browse files
committed
reset & use column schema if index is using extension dtype
1 parent d9e67c9 commit bf8f329

File tree

1 file changed

+11
-5
lines changed

1 file changed

+11
-5
lines changed

awswrangler/_data_types.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -468,11 +468,17 @@ def pyarrow_types_from_pandas( # pylint: disable=too-many-branches
468468
indexes: List[str] = []
469469
if index is True:
470470
# Get index columns
471-
# Adding indexes as columns via .reset_index() because
472-
# pa.Schema.from_pandas(.., preserve_index=True) fails with
473-
# "'Index' object has no attribute 'head'" if using extension
474-
# dtypes on pandas 1.4.x
475-
fields = pa.Schema.from_pandas(df=df[[]].reset_index(), preserve_index=False)
471+
try:
472+
fields = pa.Schema.from_pandas(df=df[[]], preserve_index=True)
473+
except AttributeError as ae:
474+
if "'Index' object has no attribute 'head'" not in str(ae):
475+
raise ae
476+
# Get index fields from a new df with only index columns
477+
# Adding indexes as columns via .reset_index() because
478+
# pa.Schema.from_pandas(.., preserve_index=True) fails with
479+
# "'Index' object has no attribute 'head'" if using extension
480+
# dtypes on pandas 1.4.x
481+
fields = pa.Schema.from_pandas(df=df.reset_index().drop(columns=cols), preserve_index=False)
476482
for field in fields:
477483
name = str(field.name)
478484
_logger.debug("Inferring PyArrow type from index: %s", name)

0 commit comments

Comments
 (0)