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

Jump to bottom

Handle pandas timestamps #958

Merged
merged 15 commits into from
Dec 1, 2022
Merged

Handle pandas timestamps #958

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
d0d5e27
Checks schema of candidate dataframe and if necessary rebuilds the sc…
hayesgb Nov 26, 2022
abb17ce
update comments
hayesgb Nov 26, 2022
e42b98c
Add test for schema provided on write_deltalake
hayesgb Nov 27, 2022
89f8dea
Linting
hayesgb Nov 27, 2022
0d726a4
Move helper function to public
hayesgb Nov 27, 2022
38b6b7b
Fix mypy
hayesgb Nov 28, 2022
c6ee10a
Linting
hayesgb Nov 30, 2022
6348ca1
Moved delta_arrow_schema_from_pandas to schema.py
hayesgb Nov 30, 2022
127abf5
Update python/deltalake/schema.py
hayesgb Nov 30, 2022
0067cd6
Update python/deltalake/schema.py
hayesgb Nov 30, 2022
07aa60b
Update python/deltalake/schema.py
hayesgb Nov 30, 2022
7f8530f
Update python/deltalake/schema.py
hayesgb Dec 1, 2022
d648f46
Update python/deltalake/schema.py
hayesgb Dec 1, 2022
1750839
fix: make pandas optional again
wjones127 Dec 1, 2022
e8fb3be
Apply suggestions from code review
wjones127 Dec 1, 2022
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
22 changes: 21 additions & 1 deletion python/deltalake/writer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,8 +132,28 @@ def write_deltalake(
:param overwrite_schema: If True, allows updating the schema of the table.
:param storage_options: options passed to the native delta filesystem. Unused if 'filesystem' is defined.
"""

if _has_pandas and isinstance(data, pd.DataFrame):
data = pa.Table.from_pandas(data)
if schema is not None:
data = pa.Table.from_panda(data, schema=schema)
wjones127 marked this conversation as resolved.
Show resolved Hide resolved
else:
_data = pa.Table.from_pandas(data)
_schema = _data.schema
schema_out = []
for _field in _schema:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would you be willing to pull this out into a helper function (_delta_arrow_schema_from_pandas)?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Really like this idea. Factored this out into a helper function, but I'd prefer it to be public rather than private. Would like to leverage this helper function elsewhere. Thoughts?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, but if you make it public, then I'd suggest putting it in schema.py.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

# partially handles https://github.com/delta-io/delta-rs/issues/686
if isinstance(_field.type, pa.lib.TimestampType):
f = pa.field(
name=_field.name,
type=pa.timestamp("us"),
nullable=_field.nullable,
metadata=_field.metadata,
)
schema_out.append(f)
else:
schema_out.append(_field)
schema = pa.schema(schema_out, metadata=_schema.metadata)
data = pa.Table.from_pandas(data, schema=schema)

if schema is None:
if isinstance(data, RecordBatchReader):
Expand Down
3 changes: 1 addition & 2 deletions python/tests/test_writer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -275,9 +275,8 @@ def test_fails_wrong_partitioning(existing_table: DeltaTable, sample_data: pa.Ta
def test_write_pandas(tmp_path: pathlib.Path, sample_data: pa.Table):
# When timestamp is converted to Pandas, it gets casted to ns resolution,
# but Delta Lake schemas only support us resolution.
sample_pandas = sample_data.to_pandas().drop(["timestamp"], axis=1)
sample_pandas = sample_data.to_pandas()
write_deltalake(str(tmp_path), sample_pandas)

delta_table = DeltaTable(str(tmp_path))
df = delta_table.to_pandas()
assert_frame_equal(df, sample_pandas)
Expand Down