{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Hi peeps,\n\nI was trying to extract data from parquet, one event_id at a time. DuckDB (https://duckdb.org/2021/06/25/querying-parquet.html) seems an interesting tool.\nHere's what I came up with:","metadata":{}},{"cell_type":"code","source":"import pandas as pd\n! pip install duckdb\nimport duckdb","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-01-25T23:01:40.158141Z","iopub.execute_input":"2023-01-25T23:01:40.158604Z","iopub.status.idle":"2023-01-25T23:01:52.558810Z","shell.execute_reply.started":"2023-01-25T23:01:40.158512Z","shell.execute_reply":"2023-01-25T23:01:52.557929Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Class to use sql for querying the data in parquet file:","metadata":{}},{"cell_type":"code","source":"class Select:\n    \"\"\"\n    Takes a path to parquet file and returns sql `select` queries results, as dataframes\n    \"\"\"\n    def __init__(self, parquet_path, columns, where='1=1'):\n        \"\"\"\n        param: parquet_path (str). The path to parquet file\n        param: columns (list or str). A list of columns to use in query, or a single column name\n        param: where (str). The sql where syntax. E.g.: 'event_id=41'\n        \"\"\"\n        assert isinstance(parquet_path, str), f\"Path to parquet must be string. Got {parquet_path, type(parquet_path)}\"\n        assert isinstance(columns, (list, str)), f\"Columns must be list or string. Got {columns, type(columns)}\"\n        assert isinstance(where, str), f\"Where for query must be string. Got {where, type(where)}\"\n        self.parquet_path = parquet_path\n        self.cols = ', '.join(columns) if isinstance(columns, list) else columns\n        self.where = where\n\n    def distinct(self):\n        query = f\"\"\"\n        SELECT DISTINCT {self.cols}\n        FROM '{self.parquet_path}'\n        \"\"\"\n        return duckdb.query(query).to_df()\n\n    def everything(self):\n        query = f\"\"\"\n        SELECT {self.cols}\n        FROM '{self.parquet_path}'\n        WHERE {self.where}\n        \"\"\"\n        return duckdb.query(query).to_df()","metadata":{"execution":{"iopub.status.busy":"2023-01-25T23:01:52.560653Z","iopub.execute_input":"2023-01-25T23:01:52.560932Z","iopub.status.idle":"2023-01-25T23:01:52.569579Z","shell.execute_reply.started":"2023-01-25T23:01:52.560907Z","shell.execute_reply":"2023-01-25T23:01:52.568231Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Feeder (generator) class which yields all data specific to an event_id, per call. Merged `azimuth` and `zenith` to output","metadata":{}},{"cell_type":"code","source":"class DataFeeder:\n    \"\"\"\n    Takes a list with train files paths, the path to metadata file and yields a dataframe with all rows, specific to each event_id, from train files.\n    \"\"\"\n    def __init__(self, paths_to_data: list, path_to_metadata: str):\n        \"\"\"\n        :param paths_to_data (list of str). The list of strings, indicating paths to train files\n        :param path_to_metadata (str). The path to metadata file\n        \"\"\"\n        self.event_ids = None\n        self.paths_to_data = paths_to_data\n        self.path_to_metadata = path_to_metadata\n\n    def feed(self):\n        for path_to_data in self.paths_to_data:\n            self.event_ids = Select(path_to_data, ['event_id']).distinct().to_numpy().squeeze()\n            for event in self.event_ids:\n                data_df = Select(\n                    parquet_path=path_to_data,\n                    columns=['event_id', 'sensor_id', 'time', 'charge', 'auxiliary'],\n                    where=f'event_id={event}'\n                ).everything()\n                meta_data_df = Select(\n                    parquet_path=self.path_to_metadata,\n                    columns=['event_id', 'azimuth', 'zenith'],\n                    where=f'event_id={event}'\n                ).everything()\n                yield pd.merge(left=data_df, right=meta_data_df, on='event_id')","metadata":{"execution":{"iopub.status.busy":"2023-01-25T23:01:52.572145Z","iopub.execute_input":"2023-01-25T23:01:52.572421Z","iopub.status.idle":"2023-01-25T23:01:52.579240Z","shell.execute_reply.started":"2023-01-25T23:01:52.572374Z","shell.execute_reply":"2023-01-25T23:01:52.578656Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Usage:","metadata":{"execution":{"iopub.status.busy":"2023-01-25T23:00:01.312730Z","iopub.execute_input":"2023-01-25T23:00:01.313221Z","iopub.status.idle":"2023-01-25T23:00:01.322819Z","shell.execute_reply.started":"2023-01-25T23:00:01.313183Z","shell.execute_reply":"2023-01-25T23:00:01.320751Z"}}},{"cell_type":"code","source":"dtf = DataFeeder(\n    ['/kaggle/input/icecube-neutrinos-in-deep-ice/train/batch_1.parquet'],\n    '/kaggle/input/icecube-neutrinos-in-deep-ice/train_meta.parquet').feed()","metadata":{"execution":{"iopub.status.busy":"2023-01-25T23:01:52.580936Z","iopub.execute_input":"2023-01-25T23:01:52.581427Z","iopub.status.idle":"2023-01-25T23:01:52.593470Z","shell.execute_reply.started":"2023-01-25T23:01:52.581402Z","shell.execute_reply":"2023-01-25T23:01:52.592454Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"next(dtf)","metadata":{"execution":{"iopub.status.busy":"2023-01-25T23:01:52.595508Z","iopub.execute_input":"2023-01-25T23:01:52.595925Z","iopub.status.idle":"2023-01-25T23:01:58.299245Z","shell.execute_reply.started":"2023-01-25T23:01:52.595896Z","shell.execute_reply":"2023-01-25T23:01:58.298262Z"},"trusted":true},"execution_count":null,"outputs":[]}]}