{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4327ef2f-39cf-4890-bcc7-372b6a4b837d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# deps:\n",
    "# pip install deltalake emmet-core pyarrow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7da9b9b3-7d71-4377-8a53-ed3084425fd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow as pa\n",
    "import pyarrow.parquet as pq\n",
    "from deltalake import DeltaTable, QueryBuilder\n",
    "from emmet.core.summary import SummaryDoc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28669cc1-f603-49d6-8577-9dcb1f032b5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_ids = [\n",
    "    \"mp-aaaaaabw\",\n",
    "    \"mp-aaacgmaf\",\n",
    "    \"mp-aaaaaaiq\",\n",
    "    \"mp-aaaaacsb\",\n",
    "    \"mp-aaaaablw\",\n",
    "    \"mp-aaaaaagb\",\n",
    "    \"mp-aaaaackw\",\n",
    "    \"mp-aaaaaazp\",\n",
    "    \"mp-aaaabhot\",\n",
    "    \"mp-aaaaabqx\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ae37e9d-6420-4d0c-8a86-c6b5734236de",
   "metadata": {},
   "outputs": [],
   "source": [
    "id_str = \",\".join(f\"'{i}'\" for i in target_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94aed126-c1f2-4f98-8250-d4570de8eebe",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_tbl = DeltaTable(\n",
    "    \"s3a://materialsproject-build/collections/summary\",\n",
    "    storage_options={\n",
    "        \"AWS_REGION\": \"us-east-1\",\n",
    "        \"AWS_SKIP_SIGNATURE\": \"true\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48612567-b50d-4150-b8a1-117331a98765",
   "metadata": {},
   "outputs": [],
   "source": [
    "qb = QueryBuilder()\n",
    "qb.register(\"summary\", summary_tbl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2c22ba4-9b51-4a8f-850f-b3b279bfe311",
   "metadata": {},
   "outputs": [],
   "source": [
    "# table name in FROM clause must match string registered to query builder above^\n",
    "# versions can be found at: https://docs.materialsproject.org/changes/database-versions\n",
    "# as of today (August 07, 2026) only most recent version is available, previous versions\n",
    "# are planned to be backfilled\n",
    "\n",
    "query_str = f\"\"\"\n",
    "    SELECT *\n",
    "    FROM   summary\n",
    "    WHERE  version='2026-04-13'\n",
    "      AND  material_id IN ({id_str})\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c56ae32d-488f-4da4-8d3d-8f62c9df5562",
   "metadata": {},
   "outputs": [],
   "source": [
    "results: pa.Table = pa.table(qb.execute(query_str).read_all())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f844f834-9abd-4db3-922a-59cb05136e27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# choice here:\n",
    "# 1) write parquet file directly using `results` table, or parquet dataset partioned on version: https://arrow.apache.org/docs/python/dataset.html#writing-datasets\n",
    "# 2) Convert to python/json and store as json, or in some persistent db system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ea03209-3ad6-4684-94b0-83ba56ae12e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# going to assume #2 for json since there are some extra steps worth illustrating ->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b303e253-b4c2-4f96-b157-c6631a9127a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "as_pydantic: list[SummaryDoc] = [SummaryDoc(**d) for d in results.to_pylist(maps_as_pydicts=\"strict\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7052481b-7229-49c7-9220-cf01c21bd033",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cast material_id to legacy format for familiarity\n",
    "for doc in as_pydantic:\n",
    "    doc.material_id = str(doc.material_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "207477bd-2ada-4db2-b706-092170421cff",
   "metadata": {},
   "outputs": [],
   "source": [
    "as_py: list[dict] = [doc.model_dump() for doc in as_pydantic]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ed93dcb-c747-4971-bf06-f2b730c0c3ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(\"materials.json\", \"w\") as f:\n",
    "#     json.dump(as_py, f)\n",
    "# etc., etc."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pipelines",
   "language": "python",
   "name": "pipelines"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}