query_data.py

from metaflow import FlowSpec, pypi, step, card, resources, Flow, S3


class QueryData(FlowSpec):

    @resources(memory=16000)
    @card
    @pypi(packages={"duckdb": "1.4.3", "pandas": "3.0.3"}, python="3.12")
    @step
    def start(self):
        import duckdb  # pylint: disable=import-error
        import pandas  # pylint: disable=import-error

        self.files = Flow("PrepareData").latest_successful_run.data.files

        with S3() as s3:
            objs = s3.get_many(self.files)
            paths = [obj.path for obj in objs]
            print(f"Found {len(paths)} parquet files")
            self.df = duckdb.query(
                """
                select * from
                read_parquet(?)
                limit 100
                """,
                params=[paths],
            ).df()
        self.next(self.end)

    @step
    def end(self):
        pass


if __name__ == "__main__":
    QueryData()