address comments

This commit is contained in:
Xi Yan 2024-10-22 15:08:57 -07:00
parent cf90ba444d
commit f131c28db7

View file

@ -48,8 +48,7 @@ class PandasDataframeDataset(BaseDataset):
self.df = None self.df = None
def __len__(self) -> int: def __len__(self) -> int:
if self.df is None: assert self.df is not None, "Dataset not loaded. Please call .load() first"
raise ValueError("Dataset not loaded. Please call .load() first")
return len(self.df) return len(self.df)
def __getitem__(self, idx): def __getitem__(self, idx):
@ -59,7 +58,7 @@ class PandasDataframeDataset(BaseDataset):
return self.df.iloc[idx].to_dict() return self.df.iloc[idx].to_dict()
def load(self) -> None: def load(self) -> None:
if self.df: if self.df is not None:
return return
# TODO: more robust support w/ data url # TODO: more robust support w/ data url
@ -106,7 +105,6 @@ class MetaReferenceDatasetioImpl(DatasetIO, DatasetsProtocolPrivate):
dataset_def: DatasetDef, dataset_def: DatasetDef,
) -> None: ) -> None:
dataset_impl = PandasDataframeDataset(dataset_def) dataset_impl = PandasDataframeDataset(dataset_def)
dataset_impl.load()
self.dataset_infos[dataset_def.identifier] = DatasetInfo( self.dataset_infos[dataset_def.identifier] = DatasetInfo(
dataset_def=dataset_def, dataset_def=dataset_def,
dataset_impl=dataset_impl, dataset_impl=dataset_impl,
@ -123,6 +121,8 @@ class MetaReferenceDatasetioImpl(DatasetIO, DatasetsProtocolPrivate):
filter_condition: Optional[str] = None, filter_condition: Optional[str] = None,
) -> PaginatedRowsResult: ) -> PaginatedRowsResult:
dataset_info = self.dataset_infos.get(dataset_id) dataset_info = self.dataset_infos.get(dataset_id)
dataset_info.dataset_impl.load()
if page_token is None: if page_token is None:
next_page_token = 0 next_page_token = 0
else: else: