Skip to content
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

Optionally allow keeping original column names in fetch_dataframe method #232

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Optionally allow keeping original column names in fetch_dataframe method #232

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
18 changes: 15 additions & 3 deletions redshift_connector/cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -504,13 +504,18 @@ def __next__(self: "Cursor") -> typing.List:
else:
raise StopIteration()

def fetch_dataframe(self: "Cursor", num: typing.Optional[int] = None) -> "pandas.DataFrame":
def fetch_dataframe(
self: "Cursor",
num: typing.Optional[int] = None,
lowercase_column_names: bool = True,
) -> "pandas.DataFrame":
"""
Fetches a user defined number of rows of a query result as a :class:`pandas.DataFrame`.

Parameters
----------
num : Optional[int] The number of rows to retrieve. If unspecified, all rows will be retrieved
lowercase_column_names : bool If set to True, column names in returend dataframe will be in lower case. Else, original column names are returned.

Returns
-------
Expand All @@ -519,11 +524,18 @@ def fetch_dataframe(self: "Cursor", num: typing.Optional[int] = None) -> "pandas
try:
import pandas
except ModuleNotFoundError:
raise ModuleNotFoundError(MISSING_MODULE_ERROR_MSG.format(module="pandas"))
raise ModuleNotFoundError(MISSING_MODULE_ERROR_MSG.format(module="pandas"))

def _handle_name(name: typing.Union[str, bytes]) -> typing.Union[str, bytes]:
"""Inner helper to optionally handle column name"""
output = name
if lowercase_column_names:
output = name.lower()
return output

columns: typing.Optional[typing.List[typing.Union[str, bytes]]] = None
try:
columns = [column[0].lower() for column in self.description]
columns = [_handle_name(column[0]) for column in self.description]
except:
warn("No row description was found. pandas dataframe will be missing column labels.", stacklevel=2)

Expand Down