Skip to content

Commit 06e99db

Browse files
authored
fix: fix list failing without order_by and local sorting (#320)
1 parent 272fc68 commit 06e99db

File tree

2 files changed

+22
-7
lines changed

2 files changed

+22
-7
lines changed

google/cloud/aiplatform/base.py

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -923,14 +923,15 @@ def _list_with_local_order(
923923
credentials=credentials,
924924
)
925925

926-
desc = "desc" in order_by
927-
order_by = order_by.replace("desc", "")
928-
order_by = order_by.split(",")
926+
if order_by:
927+
desc = "desc" in order_by
928+
order_by = order_by.replace("desc", "")
929+
order_by = order_by.split(",")
929930

930-
li.sort(
931-
key=lambda x: tuple(getattr(x, field.strip()) for field in order_by),
932-
reverse=desc,
933-
)
931+
li.sort(
932+
key=lambda x: tuple(getattr(x, field.strip()) for field in order_by),
933+
reverse=desc,
934+
)
934935

935936
return li
936937

tests/unit/aiplatform/test_datasets.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -849,6 +849,20 @@ def test_list_dataset(self, list_datasets_mock):
849849
for ds in ds_list:
850850
assert type(ds) == aiplatform.TabularDataset
851851

852+
def test_list_dataset_no_order_or_filter(self, list_datasets_mock):
853+
854+
ds_list = aiplatform.TabularDataset.list()
855+
856+
list_datasets_mock.assert_called_once_with(
857+
request={"parent": _TEST_PARENT, "filter": None}
858+
)
859+
860+
# Ensure returned list is smaller since it filtered out non-tabular datasets
861+
assert len(ds_list) < len(_TEST_DATASET_LIST)
862+
863+
for ds in ds_list:
864+
assert type(ds) == aiplatform.TabularDataset
865+
852866

853867
class TestTextDataset:
854868
def setup_method(self):

0 commit comments

Comments
 (0)