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.

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

Fix slow DataFrame creation in CogniteResource.to_pandas #1389

Merged
merged 12 commits into from
Oct 4, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
6 changes: 2 additions & 4 deletions cognite/client/data_classes/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,10 +164,8 @@ def to_pandas(
dumped.update(dumped.pop(key))
else:
raise AssertionError(f"Could not expand attribute '{key}'")
df = pd.DataFrame(columns=["value"])
for name, value in dumped.items():
df.loc[name] = [value]
return df

return pd.Series(dumped).to_frame(name="value")

def _repr_html_(self) -> str:
return notebook_display_with_fallback(self)
Expand Down
35 changes: 35 additions & 0 deletions tests/tests_unit/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,41 @@ def test_use_method_which_requires_cognite_client__client_not_set(self):
with pytest.raises(CogniteMissingClientError):
mr.use()

@pytest.mark.dsl
def test_to_pandas_method(self):
import pandas as pd

from cognite.client.data_classes import Asset, Label

asset = Asset(
external_id="test-1",
name="test 1",
parent_external_id="parent-test-1",
description="A test asset",
data_set_id=123,
labels=[Label(external_id="ROTATING_EQUIPMENT", name="Rotating equipment")],
)

result_df = asset.to_pandas()

data = {
"value": [
"test-1",
"test 1",
"parent-test-1",
"A test asset",
123,
[{"externalId": "ROTATING_EQUIPMENT", "name": "Rotating equipment"}],
]
}

index_labels = ["external_id", "name", "parent_external_id", "description", "data_set_id", "labels"]

expected_df = pd.DataFrame(data, index=index_labels)

# Assert that the resultant DataFrame is equal to the expected DataFrame
pd.testing.assert_frame_equal(result_df, expected_df)


class TestCogniteFilter:
def test_dump(self):
Expand Down