importpandasaspdfrommetaflowimportFlowSpec,stepfromzdatasetsimportDataset,Modefromzdatasets.metaflowimportDatasetParameterfromzdatasets.pluginsimportBatchOptions# Can also invoke from CLI:# > python zdatasets/tutorials/0_hello_dataset_flow.py run \# --hello_dataset '{"name": "HelloDataset", "mode": "READ_WRITE", \# "options": {"type": "BatchOptions", "partition_by": "region"}}'classHelloDatasetFlow(FlowSpec):hello_dataset=DatasetParameter("hello_dataset",default=Dataset("HelloDataset",mode=Mode.READ_WRITE,options=BatchOptions(partition_by="region")), )@stepdefstart(self):df=pd.DataFrame({"region": ["A","A","A","B","B","B"],"zpid": [1,2,3,4,5,6]})print("saving data_frame:\n",df.to_string(index=False))# Example of writing to a datasetself.hello_dataset.write(df)# save this as an output datasetself.output_dataset=self.hello_datasetself.next(self.end)@stepdefend(self):print(f"I have dataset\n{self.output_dataset=}")# output_dataset to_pandas(partitions=dict(region="A")) onlydf:pd.DataFrame=self.output_dataset.to_pandas(partitions=dict(region="A"))print('self.output_dataset.to_pandas(partitions=dict(region="A")):')print(df.to_string(index=False))if__name__=="__main__":HelloDatasetFlow()
