{"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":"# XLA Tile Preprocessing with Polars\n\nThis notebook is intended to show to prepare the Tile XLA data into a more manageable format (parquet), as well as how to use Polars to efficiently manipulate and process parquet files in a lazy way.\nThis notebook is a two-parter.\n\n* Part 1: Processing our Data into Parquet\n  A simple script that parses the tile XLA data into parquet files in 3 tables - edge, node, and config, with the appropriate column names. It's nice to have this on disk.\n\n* Part 2: Lazy Processing with Polars\n  An example of how to set up a lazy pipeline to process our parquet data into a usable format.","metadata":{}},{"cell_type":"code","source":"import os\nimport polars as pl\nimport numpy as np\nfrom typing import Any\nfrom dataclasses import dataclass\n\nfrom tqdm.auto import tqdm\nfrom polars.type_aliases import IntoExpr\nimport networkx as nx","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-10-06T13:34:44.051569Z","iopub.execute_input":"2023-10-06T13:34:44.051929Z","iopub.status.idle":"2023-10-06T13:34:44.141291Z","shell.execute_reply.started":"2023-10-06T13:34:44.051905Z","shell.execute_reply":"2023-10-06T13:34:44.140438Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Some index to feature name maps are defined below. I prefer working with column names when possible.","metadata":{}},{"cell_type":"code","source":"NODE_FEATURE_MAP = {\n    0: \"is_root\",  # - whether this node is the output,\n    1: \"element_size_in_bits\",  # - deprecated, always 0,\n    # 2\"–\"20: One hot vector of shape_element_type.,\n    2: \"shape_element_type_is_invalid_type\",\n    3: \"shape_element_type_is_pred\",\n    4: \"shape_element_type_is_s8\",\n    5: \"shape_element_type_is_s16\",\n    6: \"shape_element_type_is_s32\",\n    7: \"shape_element_type_is_s64\",\n    8: \"shape_element_type_is_u8\",\n    9: \"shape_element_type_is_u16\",\n    10: \"shape_element_type_is_u32\",\n    11: \"shape_element_type_is_u64\",\n    12: \"shape_element_type_is_f16\",\n    13: \"shape_element_type_is_f32\",\n    14: \"shape_element_type_is_f64\",\n    15: \"shape_element_type_is_bf16\",\n    16: \"shape_element_type_is_c64\",\n    17: \"shape_element_type_is_c128\",\n    18: \"shape_element_type_is_tuple\",\n    19: \"shape_element_type_is_opaque_type\",\n    20: \"shape_element_type_is_token\",\n    # \"21\"–28: Size (number of elements) for each dimension, or an upper bound on the size if the dimension is dynamic.  In XLA, dimensions are numbered from 0 to N-1 for an N-dimensional array. The first element of 'shape_dimensions' is the size of dimension 0, the second element is the size of dimension 1, and so forth.  Empty list indicates a scalar.,\n    21: \"shape_dimensions_0\",\n    22: \"shape_dimensions_1\",\n    23: \"shape_dimensions_2\",\n    24: \"shape_dimensions_3\",\n    25: \"shape_dimensions_4\",\n    26: \"shape_dimensions_5\",\n    27: \"shape_dimensions_sum\",\n    28: \"shape_dimensions_product\",\n    29: \"shape_tuple_shapes_size\",  # - for tuples only, the shapes of constituent shapes in the tuple sequence,\n    30: \"parameter_number\",  # = K - indicating that is is the Kth parameter to the computation, only for Parameter operation,\n    # \"31\"–36: Dimensions present for some operations that require reshaping or broadcasting, including Reshape, Reduce, ReduceWindow, and Reverse.,\n    31: \"dimensions_0\",\n    32: \"dimensions_1\",\n    33: \"dimensions_2\",\n    34: \"dimensions_3\",\n    35: \"dimensions_4\",\n    36: \"dimensions_5\",\n    # \"37\"–92: Windowing information in an operation such as convolution. The window is moved across a base area and for each position of the window a computation is performed.,\n    37: \"window_size_0\",\n    38: \"window_size_1\",\n    39: \"window_size_2\",\n    40: \"window_size_3\",\n    41: \"window_size_4\",\n    42: \"window_size_5\",\n    43: \"window_size_sum\",\n    44: \"window_size_product\",\n    45: \"window_stride_0\",\n    46: \"window_stride_1\",\n    47: \"window_stride_2\",\n    48: \"window_stride_3\",\n    49: \"window_stride_4\",\n    50: \"window_stride_5\",\n    51: \"window_stride_sum\",\n    52: \"window_stride_product\",\n    53: \"window_padding_low_0\",\n    54: \"window_padding_low_1\",\n    55: \"window_padding_low_2\",\n    56: \"window_padding_low_3\",\n    57: \"window_padding_low_4\",\n    58: \"window_padding_low_5\",\n    59: \"window_padding_low_sum\",\n    60: \"window_padding_low_product\",\n    61: \"window_padding_high_0\",\n    62: \"window_padding_high_1\",\n    63: \"window_padding_high_2\",\n    64: \"window_padding_high_3\",\n    65: \"window_padding_high_4\",\n    66: \"window_padding_high_5\",\n    67: \"window_padding_high_sum\",\n    68: \"window_padding_high_product\",\n    # \"69\"–76: Dilation factor of the sliding window. A dilation factor of 1 means no dilation. window_dilation - 1 no-op entries (\"holes\") are implicitly placed between each kernel element.,\n    69: \"window_window_dilation_0\",\n    70: \"window_window_dilation_1\",\n    71: \"window_window_dilation_2\",\n    72: \"window_window_dilation_3\",\n    73: \"window_window_dilation_4\",\n    74: \"window_window_dilation_5\",\n    75: \"window_window_dilation_sum\",\n    76: \"window_window_dilation_product\",\n    # \"77\"-84: Dilation factor of the base area. A dilation factor of 1 means no dilation. base_dilation - 1 no-op entries (\"holes\") are implicitly placed between each base area element.,\n    77: \"window_base_dilation_0\",\n    78: \"window_base_dilation_1\",\n    79: \"window_base_dilation_2\",\n    80: \"window_base_dilation_3\",\n    81: \"window_base_dilation_4\",\n    82: \"window_base_dilation_5\",\n    83: \"window_base_dilation_sum\",\n    84: \"window_base_dilation_product\",\n    # \"85\"-92: Window reversal means that this dimension was logically reversed before the operation.,\n    85: \"window_window_reversal_0\",\n    86: \"window_window_reversal_1\",\n    87: \"window_window_reversal_2\",\n    88: \"window_window_reversal_3\",\n    89: \"window_window_reversal_4\",\n    90: \"window_window_reversal_5\",\n    91: \"window_window_reversal_true_count\",\n    92: \"window_window_reversal_false_count\",\n    # \"93\"–106: The dimension numbers used for a convolution.,\n    93: \"convolution_dim_numbers_input_batch_dim\",  # - the dimension number that represents batch in the input,\n    94: \"convolution_dim_numbers_input_feature_dim\",  # - the dimension number that represents features in the input,\n    # \"95\"–98: Dimension numbers for the spatial dimensions that the window moves through in the input.,\n    95: \"convolution_dim_numbers_input_spatial_dims_0\",\n    96: \"convolution_dim_numbers_input_spatial_dims_1\",\n    97: \"convolution_dim_numbers_input_spatial_dims_2\",\n    98: \"convolution_dim_numbers_input_spatial_dims_3\",\n    99: \"convolution_dim_numbers_kernel_input_feature_dim\",  # - the dimension number that represents input features in the convolutional kernel (rhs),\n    100: \"convolution_dim_numbers_kernel_output_feature_dim\",  # - the dimension number that represents output features in the convolutional kernel (rhs),\n    # \"101\"-104: Dimension numbers for the spatial dimensions that the window moves through in the kernel (rhs). window.strides(0) is the stride in the kernel_spatial_dimensions(0) dimension.,\n    101: \"convolution_dim_numbers_kernel_spatial_dims_0\",\n    102: \"convolution_dim_numbers_kernel_spatial_dims_1\",\n    103: \"convolution_dim_numbers_kernel_spatial_dims_2\",\n    104: \"convolution_dim_numbers_kernel_spatial_dims_3\",\n    105: \"convolution_dim_numbers_output_batch_dim\",  # - the dimension number that represents batch in the output,\n    106: \"convolution_dim_numbers_output_feature_dim\",  # - the dimension number that represents features in the output,\n    107: \"feature_group_count\",  # - the number of feature groups, used for a convolution. Must be a divisor of the input feature dimension and output feature dimension. If not specified, it will use a default value of 1.,\n    108: \"batch_group_count\",  # - the number of batch groups, used for a convolution.,\n    # \"109\"–120: [begin/start, end/limit) index range and stride for a slice operation.,\n    109: \"slice_dims_start_0\",\n    110: \"slice_dims_start_1\",\n    111: \"slice_dims_start_sum\",\n    112: \"slice_dims_start_product\",\n    113: \"slice_dims_stride_0\",\n    114: \"slice_dims_stride_1\",\n    115: \"slice_dims_stride_sum\",\n    116: \"slice_dims_stride_product\",\n    117: \"slice_dims_limit_0\",\n    118: \"slice_dims_limit_1\",\n    119: \"slice_dims_limit_sum\",\n    120: \"slice_dims_limit_product\",\n    # \"121\" - 124: [start, start + size) range size for a dynamic slice ('start' is specified dynamically in the second operand of the operation).,\n    121: \"dynamic_slice_sizes_0\",\n    122: \"dynamic_slice_sizes_1\",\n    123: \"dynamic_slice_sizes_sum\",\n    124: \"dynamic_slice_sizes_product\",\n    # \"125\"–132: Padding configuration that describes the edge padding of a pad operation.,\n    125: \"padding_config_edge_padding_low_0\",\n    126: \"padding_config_edge_padding_low_1\",\n    127: \"padding_config_edge_padding_low_sum\",\n    128: \"padding_config_edge_padding_low_product\",\n    129: \"padding_config_edge_padding_high_0\",\n    130: \"padding_config_edge_padding_high_1\",\n    131: \"padding_config_edge_padding_high_sum\",\n    132: \"padding_config_edge_padding_high_product\",\n    133: \"is_stable\",  # - whether this Sort operation should be stable,\n    # \"134\"–139: Physical layout used to pack the tensor shape.,\n    134: \"layout_minor_to_major_0\",\n    135: \"layout_minor_to_major_1\",\n    136: \"layout_minor_to_major_2\",\n    137: \"layout_minor_to_major_3\",\n    138: \"layout_minor_to_major_4\",\n    139: \"layout_minor_to_major_5\",\n}\n\nCONFIG_FEATURE_MAP = {\n    # 0–7: Tile sizes of the convolution kernel, only for a convolution operation.,\n    0: \"kernel_bounds_0\",\n    1: \"kernel_bounds_1\",\n    2: \"kernel_bounds_2\",\n    3: \"kernel_bounds_3\",\n    4: \"kernel_bounds_4\",\n    5: \"kernel_bounds_5\",\n    6: \"kernel_bounds_sum\",\n    7: \"kernel_bounds_product\",\n    # 8\"–\"15: Output tile sizes.,\n    8: \"output_bounds_0\",\n    9: \"output_bounds_1\",\n    10: \"output_bounds_2\",\n    11: \"output_bounds_3\",\n    12: \"output_bounds_4\",\n    13: \"output_bounds_5\",\n    14: \"output_bounds_sum\",\n    15: \"output_bounds_product\",\n    # \"16\"-23: Input tile sizes.,\n    16: \"input_bounds_0\",\n    17: \"input_bounds_1\",\n    18: \"input_bounds_2\",\n    19: \"input_bounds_3\",\n    20: \"input_bounds_4\",\n    21: \"input_bounds_5\",\n    22: \"input_bounds_sum\",\n    23: \"input_bounds_product\",\n}\n\nOPCODE_MAP = {\n    1: \"abs\",\n    2: \"add\",\n    3: \"add-dependency\",\n    4: \"after-all\",\n    5: \"all-reduce\",\n    6: \"all-to-all\",\n    7: \"atan2\",\n    8: \"batch-norm-grad\",\n    9: \"batch-norm-inference\",\n    10: \"batch-norm-training\",\n    11: \"bitcast\",\n    12: \"bitcast-convert\",\n    13: \"broadcast\",\n    14: \"call\",\n    15: \"ceil\",\n    16: \"cholesky\",\n    17: \"clamp\",\n    18: \"collective-permute\",\n    19: \"count-leading-zeros\",\n    20: \"compare\",\n    21: \"complex\",\n    22: \"concatenate\",\n    23: \"conditional\",\n    24: \"constant\",\n    25: \"convert\",\n    26: \"convolution\",\n    27: \"copy\",\n    28: \"copy-done\",\n    29: \"copy-start\",\n    30: \"cosine\",\n    31: \"custom-call\",\n    32: \"divide\",\n    33: \"domain\",\n    34: \"dot\",\n    35: \"dynamic-slice\",\n    36: \"dynamic-update-slice\",\n    37: \"exponential\",\n    38: \"exponential-minus-one\",\n    39: \"fft\",\n    40: \"floor\",\n    41: \"fusion\",\n    42: \"gather\",\n    43: \"get-dimension-size\",\n    44: \"set-dimension-size\",\n    45: \"get-tuple-element\",\n    46: \"imag\",\n    47: \"infeed\",\n    48: \"iota\",\n    49: \"is-finite\",\n    50: \"log\",\n    51: \"log-plus-one\",\n    52: \"and\",\n    53: \"not\",\n    54: \"or\",\n    55: \"xor\",\n    56: \"map\",\n    57: \"maximum\",\n    58: \"minimum\",\n    59: \"multiply\",\n    60: \"negate\",\n    61: \"outfeed\",\n    62: \"pad\",\n    63: \"parameter\",\n    64: \"partition-id\",\n    65: \"popcnt\",\n    66: \"power\",\n    67: \"real\",\n    68: \"recv\",\n    69: \"recv-done\",\n    70: \"reduce\",\n    71: \"reduce-precision\",\n    72: \"reduce-window\",\n    73: \"remainder\",\n    74: \"replica-id\",\n    75: \"reshape\",\n    76: \"reverse\",\n    77: \"rng\",\n    78: \"rng-get-and-update-state\",\n    79: \"rng-bit-generator\",\n    80: \"round-nearest-afz\",\n    81: \"rsqrt\",\n    82: \"scatter\",\n    83: \"select\",\n    84: \"select-and-scatter\",\n    85: \"send\",\n    86: \"send-done\",\n    87: \"shift-left\",\n    88: \"shift-right-arithmetic\",\n    89: \"shift-right-logical\",\n    90: \"sign\",\n    91: \"sine\",\n    92: \"slice\",\n    93: \"sort\",\n    94: \"sqrt\",\n    95: \"subtract\",\n    96: \"tanh\",\n    97: \"trace\",\n    98: \"transpose\",\n    99: \"triangular-solve\",\n    100: \"tuple\",\n    101: \"tuple-select\",\n    102: \"while\",\n    103: \"cbrt\",\n    104: \"all-gather\",\n    105: \"collective-permute-start\",\n    106: \"collective-permute-done\",\n    107: \"logistic\",\n    108: \"dynamic-reshape\",\n    109: \"all-reduce-start\",\n    110: \"all-reduce-done\",\n    111: \"reduce-scatter\",\n    112: \"all-gather-start\",\n    113: \"all-gather-done\",\n    114: \"opt-barrier\",\n    115: \"async-start\",\n    116: \"async-update\",\n    117: \"async-done\",\n    118: \"round-nearest-even\",\n    119: \"stochastic-convert\",\n    120: \"tan\",\n}","metadata":{"execution":{"iopub.status.busy":"2023-10-06T13:21:19.201276Z","iopub.execute_input":"2023-10-06T13:21:19.201659Z","iopub.status.idle":"2023-10-06T13:21:19.226378Z","shell.execute_reply.started":"2023-10-06T13:21:19.201630Z","shell.execute_reply":"2023-10-06T13:21:19.225225Z"},"_kg_hide-input":false,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Below is a utility class that coerces the numpy data from the Tile XLA into parquet files. You don't necessarily need to do this - you can modify this to just construct the dataframes and return them in memory instead.\nBut the processing isn't super quick, so it's nice to have them on disk.","metadata":{}},{"cell_type":"code","source":"class ParquetWriter:\n    \"\"\"Process the npz files into parquet files for easier work\"\"\"\n\n    def __init__(self, source_dir: str, output_dir: str) -> None:\n        self.source_dir = source_dir\n        self.output_dir = output_dir\n\n    def process_file(self, file: str) -> None:\n        os.makedirs(self.output_dir, exist_ok=True)\n\n        d: dict[str, Any] = dict(np.load(os.path.join(self.source_dir, file)))\n        d[\"file\"] = file\n\n        self.process_node_data(d)\n        self.process_config_data(d)\n        self.process_edge_data(d)\n\n    def process_node_data(self, data: dict[str, Any]) -> None:\n        \"\"\"Process the node features\"\"\"\n        node_feat = {\n            NODE_FEATURE_MAP[j]: data[\"node_feat\"][:, j]\n            for j in range(data[\"node_feat\"].shape[1])\n        }\n\n        opcodes = [OPCODE_MAP[opc] for opc in data[\"node_opcode\"]]  # type: ignore\n        node_df = pl.DataFrame(\n            {\n                **node_feat,\n                \"node_opcode\": opcodes,\n                \"node_id\": np.arange(0, len(opcodes)),\n                \"file_id\": data[\"file\"],\n            }\n        )\n\n        node_df.write_parquet(self._make_node_filepath(data[\"file\"]))\n\n    def process_config_data(self, data: dict[str, Any]) -> None:\n        config_feat = {\n            CONFIG_FEATURE_MAP[j]: data[\"config_feat\"][:, j]\n            for j in range(data[\"config_feat\"].shape[1])\n        }\n\n        config_df = pl.DataFrame(\n            {\n                **config_feat,\n                \"config_runtime\": data[\"config_runtime\"],\n                \"config_runtime_normalizers\": data[\"config_runtime_normalizers\"],\n                \"file_id\": data[\"file\"],\n            }\n        )\n        config_df.write_parquet((self._make_config_filepath(data[\"file\"])))\n\n    def process_edge_data(self, data: dict[str, Any]) -> None:\n        edge_df = pl.DataFrame(\n            {\n                \"from\": data[\"edge_index\"][:, 0],\n                \"to\": data[\"edge_index\"][:, 1],\n                \"file_id\": data[\"file\"],\n            }\n        )\n\n        edge_df.write_parquet(self._make_edge_filepath(data[\"file\"]))\n\n    def _make_node_filepath(self, filename: str) -> str:\n        os.makedirs(os.path.join(self.output_dir, \"node\"), exist_ok=True)\n        return os.path.join(self.output_dir, \"node\", f\"{filename}.parquet\")\n\n    def _make_config_filepath(self, filename: str) -> str:\n        os.makedirs(os.path.join(self.output_dir, \"config\"), exist_ok=True)\n        return os.path.join(self.output_dir, \"config\", f\"{filename}.parquet\")\n\n    def _make_edge_filepath(self, filename: str) -> str:\n        os.makedirs(os.path.join(self.output_dir, \"edge\"), exist_ok=True)\n        return os.path.join(self.output_dir, \"edge\", f\"{filename}.parquet\")","metadata":{"execution":{"iopub.status.busy":"2023-10-06T13:21:20.885037Z","iopub.execute_input":"2023-10-06T13:21:20.885411Z","iopub.status.idle":"2023-10-06T13:21:20.898792Z","shell.execute_reply.started":"2023-10-06T13:21:20.885381Z","shell.execute_reply":"2023-10-06T13:21:20.897919Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now we just process each file into it's own `.parquet` file. There's no need to shove them all in one table, we can join them efficiently if they share a schema (which they do)","metadata":{}},{"cell_type":"code","source":"from concurrent.futures import ThreadPoolExecutor\n\n\n@dataclass\nclass WorkerArgs:\n    src_path: str\n    outpath: str\n    files: list[str]\n\n\ndef process_files(args: WorkerArgs) -> None:\n    writer = ParquetWriter(args.src_path, args.outpath)\n\n    for file in tqdm(args.files):\n        writer.process_file(file)\n\n\ndef chunk_files(files: list[str], num_chunks: int):\n    chunk_size = len(files) // num_chunks\n    return [files[i : i + chunk_size] for i in range(0, len(files), chunk_size)]\n\n\nSRC_DIRECTORY = \"/kaggle/input/predict-ai-model-runtime/npz_all/npz/tile/xla\"\nOUTPUT_DIRECTORY = \"parquet\"\nCHUNKS = 32\nMAX_WORKERS = 8\n\nsplits = [\"train\",\"valid\",\"test\"]\n\nfor split in splits:\n    src_path = os.path.join(SRC_DIRECTORY, split)\n    outpath = os.path.join(OUTPUT_DIRECTORY, split)\n    files = os.listdir(src_path)\n\n    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:\n        pool.map(\n            process_files,\n            [\n                WorkerArgs(src_path, outpath, chunk)\n                for chunk in chunk_files(files, CHUNKS)\n            ],\n        )","metadata":{"execution":{"iopub.status.busy":"2023-10-06T14:09:13.421081Z","iopub.execute_input":"2023-10-06T14:09:13.421572Z","iopub.status.idle":"2023-10-06T14:09:51.006328Z","shell.execute_reply.started":"2023-10-06T14:09:13.421502Z","shell.execute_reply":"2023-10-06T14:09:51.002335Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Part 2: Scanning and processing the Parquet Data (Lazily) with Polars\nOne of the things I love most about Polars is the lazy processing. Below is a very simple but realistic example of how to use the data we just created.\n\nWe're going to:\n1. Sum all Node Features\n2. Add a few basic graph-level features\n3. Join with configurations (restricted to X per file)","metadata":{}},{"cell_type":"code","source":"# Some utility functions\n\ndef compute_graph_level_features(df: pl.DataFrame) -> pl.DataFrame:\n    \"\"\"Compute some simple graph-level statistics from an edge list.\n    Assumes that the edge list is in the format of edge.parquet files.\"\"\"\n    # Convert to tuples\n    to = df[\"to\"].cast(pl.Int32).to_list()\n    frm = df[\"from\"].cast(pl.Int32).to_list()\n\n    edgelist = [(t, f) for t, f in zip(to, frm)]\n\n    graph = nx.DiGraph(edgelist)\n    average_degree = np.mean([d for _, d in graph.degree()])\n    average_clustering = nx.average_clustering(graph)\n    longest_path = nx.dag_longest_path_length(graph)\n\n    return pl.DataFrame(\n        {\n            \"file_id\": df[\"file_id\"][0],\n            \"average_degree\": average_degree,\n            \"average_clustering\": average_clustering,\n            \"longest_path\": longest_path,\n        },\n    )\n\n\ndef sample(df: pl.DataFrame, max_configs: int = 500) -> pl.DataFrame:\n    \"\"\"A very simplistic way to trim the max number of configs per file.\"\"\"\n    return df.sample(n=min(max_configs, df.height))\n","metadata":{"execution":{"iopub.status.busy":"2023-10-06T13:35:31.462217Z","iopub.execute_input":"2023-10-06T13:35:31.462655Z","iopub.status.idle":"2023-10-06T13:35:31.471598Z","shell.execute_reply.started":"2023-10-06T13:35:31.462622Z","shell.execute_reply":"2023-10-06T13:35:31.469962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\nNow the fun bit - use polars lazy scanning to develop a computation graph. We can lazily set up our pipeline, then execute it in a memory efficient way later.","metadata":{}},{"cell_type":"code","source":"dfs = {}\nMAX_CONFIGS_PER_FILE = 500\n\nfor split in [\"train\",\"valid\",\"test\"]:\n    node_df = pl.scan_parquet(f\"{OUTPUT_DIRECTORY}/{split}/node/*.parquet\", low_memory=True)\n    config_df = pl.scan_parquet(\n        f\"{OUTPUT_DIRECTORY}/{split}/config/*.parquet\", low_memory=True\n    )\n    edge_df = pl.scan_parquet(f\"{OUTPUT_DIRECTORY}/{split}/edge/*.parquet\", low_memory=True)\n\n    # These columns are all 0 when you sum the entire train node dataset, so we can remove them\n    # as we have no instances of them in the dataset\n    # If they appear in the test set, we'll have to either ignore them or figure out a way\n    ohe_cols: list[IntoExpr] = []\n\n    for key in NODE_FEATURE_MAP.values():  # The str\n        col = (\n            (pl.col(\"node_opcode\") == key).alias(f\"opcode_{key}\")\n        )\n        ohe_cols.append(col)\n\n    processed = node_df.with_columns(ohe_cols)\n\n    # Just sum up all columns over node id\n    aggregations = (\n        processed.select(pl.exclude(\"node_opcode\", \"node_id\")).group_by(\"file_id\").sum()\n    )\n\n    graph_features = edge_df.group_by(\"file_id\").map_groups(\n        compute_graph_level_features,\n        schema={\n            \"file_id\": pl.Utf8,\n            \"average_degree\": pl.Float32,\n            \"average_clustering\": pl.Float32,\n            \"longest_path\": pl.Float32,\n        },\n    )\n\n    processed_config = config_df.with_columns(\n        (pl.col(\"config_runtime\") / pl.col(\"config_runtime_normalizers\")).alias(\"label\")\n    ).drop(\"config_runtime\", \"config_runtime_normalizers\")\n\n\n    # Restrict number of configs per file randomly. Comment this out if you want all configs\n    # Just make sure you know how to process it efficiently as it'll take a lot of memory\n    processed_config = processed_config.group_by(\"file_id\").map_groups(\n        lambda x: sample(x, MAX_CONFIGS_PER_FILE), schema=None\n    )\n\n    # There are _a crazy amount_ of permutations for some files, so we'll randomly\n    # sample no more than X per file\n\n    processed_data = aggregations.join(\n        graph_features, on=\"file_id\", how=\"inner\", validate=\"1:1\"\n    ).join(processed_config, on=\"file_id\", how=\"inner\", validate=\"1:m\")\n\n    dfs[split] = processed_data","metadata":{"execution":{"iopub.status.busy":"2023-10-06T13:36:17.522541Z","iopub.execute_input":"2023-10-06T13:36:17.522912Z","iopub.status.idle":"2023-10-06T13:36:20.695672Z","shell.execute_reply.started":"2023-10-06T13:36:17.522883Z","shell.execute_reply":"2023-10-06T13:36:20.694664Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Behold! We've defined our whole pipeline but currently have virtually no memory used. Now we can fetch, collect, or continue to transform the data as we see fit.\n\nIf we want to execute the pipeline, we can call either `fetch` (if we just want to fetch N rows) or `collect` (if we want all of them). This will execute the query plan that polars sets up, which is quite efficient.","metadata":{}},{"cell_type":"code","source":"dfs[\"train\"].fetch(100).head()","metadata":{"execution":{"iopub.status.busy":"2023-10-06T13:47:25.832777Z","iopub.execute_input":"2023-10-06T13:47:25.833188Z","iopub.status.idle":"2023-10-06T13:47:39.985473Z","shell.execute_reply.started":"2023-10-06T13:47:25.833156Z","shell.execute_reply":"2023-10-06T13:47:39.984201Z"},"trusted":true},"execution_count":null,"outputs":[]}]}