{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":84493,"databundleVersionId":11305158,"sourceType":"competition"}],"dockerImageVersionId":31012,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import gc\nimport polars as pl\n\ntrain = pl.read_parquet(f\"/kaggle/input/jane-street-real-time-market-data-forecasting/lags.parquet/date_id=0/part-0.parquet\")\ndisplay(train.head(20))\n\nlag_cols_original = [\"date_id\", \"symbol_id\"] + [f\"responder_{idx}\" for idx in range(9)]\nlag_cols_rename = { f\"responder_{idx}\" : f\"responder_{idx}_lag_1\" for idx in range(9)}\n\ndef add_lags(df):\n    lags = df.select(pl.col(lag_cols_original))\n    lags = lags.rename(lag_cols_rename)\n    lags = lags.with_columns(date_id = pl.col('date_id') + 1,)  # lagged by 1 day\n    lags = lags.group_by([\"date_id\", \"symbol_id\"], maintain_order=True).last()  # pick up last record of previous date\n    return df.join(lags, on=[\"date_id\", \"symbol_id\"], how=\"left\")\n\ndf_train, df_valid = None, None\nfor i in range(0,10):\n    df = pl.read_parquet(f\"/kaggle/input/jane-street-real-time-market-data-forecasting/train.parquet/partition_id={i}/part-0.parquet\")\n    df = df.with_columns(\n        (pl.col(\"responder_6\") * 2).cast(pl.Int32).alias(\"label\"),        \n        pl.col('date_id').cast(pl.Int64),\n        pl.col('time_id').cast(pl.Int64),\n    )\n    time_id_max = df['time_id'].max()\n    print(df['date_id'].min(), df['date_id'].max(), df['time_id'].min(), time_id_max,)\n    training_cutoff = time_id_max - int(time_id_max // 5)\n    print(df.shape, training_cutoff)\n\n    train = df.filter(pl.col('time_id') < training_cutoff)\n    valid = df.filter(pl.col('time_id') >= training_cutoff)\n    if df_train is None:\n        df_train = train\n        df_valid = valid\n    else:\n        df_train = df_train.vstack(train)\n        df_valid = df_valid.vstack(valid)\n    del train, valid\n    _ = gc.collect()\n\ndf_train = add_lags(df_train)\ndf_valid = add_lags(df_valid)\n\nprint(df_train.shape, df_valid.shape)\nprint(df_train.columns)\ndf_train.write_parquet(\"./train.parquet\")\ndf_valid.write_parquet(\"./valid.parquet\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-11T17:13:04.513284Z","iopub.execute_input":"2025-04-11T17:13:04.513572Z","iopub.status.idle":"2025-04-11T17:15:47.630474Z","shell.execute_reply.started":"2025-04-11T17:13:04.513545Z","shell.execute_reply":"2025-04-11T17:15:47.628183Z"}},"outputs":[],"execution_count":null}]}