{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3QiMeqVqYaU3",
   "metadata": {
    "id": "3QiMeqVqYaU3"
   },
   "source": [
    "# Kumo Relational Quickstart with the NVIDIA Kumo Relational Client\n",
    "\n",
    "**Kumo Relational** is a relational foundation model for\n",
    "machine learning on enterprise data. With your relational data and a\n",
    "Predictive Query Language (PQL) query, you can generate predictions without\n",
    "training a task-specific model.\n",
    "\n",
    "This quickstart uses four tables from Microsoft's **AdventureWorks** OLTP\n",
    "sample database: customers, products, sales-order headers, and sales-order\n",
    "details. Predictions are submitted to a deployed Kumo Relational NVIDIA NIM through\n",
    "the NVIDIA Kumo Relational Client.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "HKJCzrCTYaU3",
   "metadata": {
    "id": "HKJCzrCTYaU3"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "Kumo Relational is grounded in three key world views:\n",
    "\n",
    "> **1. Enterprise data is a graph.**\n",
    "\n",
    "Enterprise data is a graph where tables are connected by keys. In this notebook, `sales_order_headers` connects customers to orders, while `sales_order_details` connects each order to its products. Together, these relationships give Kumo Relational the business context surrounding every transaction.\n",
    "\n",
    "Once we structure enterprise data this way, we can apply pre-trained Relational Graph Transformers to extract insights and patterns.\n",
    "\n",
    "> **2. With timestamps, we place events on a timeline.**\n",
    "\n",
    "By placing events on a timeline, we unlock the ability to model how things evolve. We can select any point in time and predict what is likely to happen next from the sequence and patterns in the historical data.\n",
    "\n",
    "> **3. Machine learning tasks can be described via predictive queries.**\n",
    "\n",
    "Major machine learning tasks—including regression, classification, forecasting, and recommendation—can be expressed using *Predictive Query Language (PQL)*. A PQL query identifies the outcome to predict, the entity to predict it for, and, for temporal tasks, the future prediction window.\n",
    "\n",
    "If you know SQL, PQL should feel familiar. Learn more in the [PQL documentation](https://docs.nvidia.com/sdgm/reference/pq-reference-overview).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "Enucu67aEPZn",
   "metadata": {
    "id": "Enucu67aEPZn"
   },
   "source": [
    "**Let's get started!**\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "psushpJFYaU4",
   "metadata": {
    "id": "psushpJFYaU4"
   },
   "source": [
    "## Step 1. Install the NVIDIA Kumo Relational Client\n",
    "\n",
    "Install the SDK with the Kumo Relational extra. The Kumo Relational release wheels support\n",
    "Python 3.10, 3.11, 3.12, and 3.13.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ro8qKwXAYaU4",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "collapsed": true,
    "id": "ro8qKwXAYaU4",
    "outputId": "f6711111-f3ab-4c81-ddfc-5d2f8de483a2",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%pip install \"kumo-relational-client[relational]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b73a9a4-2bd6-438b-9b7b-91583cf2d85d",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 314
    },
    "id": "3b73a9a4-2bd6-438b-9b7b-91583cf2d85d",
    "outputId": "cc4914da-ca12-43ba-f2f0-7c6afd0541b6"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "import pandas as pd\n",
    "from kumo_relational_client import RelationalClient, relational"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cG9-xYdHEhH3",
   "metadata": {
    "id": "cG9-xYdHEhH3"
   },
   "source": [
    "The supported Kumo Relational object namespace is\n",
    "`kumo_relational_client.relational`. Build graphs through that namespace and submit\n",
    "predictions through `RelationalClient`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "xr7eIPGuYaU4",
   "metadata": {
    "id": "xr7eIPGuYaU4"
   },
   "source": [
    "## Step 2. Configure the Kumo Relational NIM endpoint\n",
    "\n",
    "Deploy Kumo Relational as an NVIDIA NIM before running this notebook. See\n",
    "[Deploy, Install, and Connect](https://docs.nvidia.com/sdgm/rfm/sdk-getting-started)\n",
    "for deployment instructions.\n",
    "\n",
    "Set `KUMO_RELATIONAL_API_ENDPOINT` to the reachable NIM endpoint. The default below uses a\n",
    "local container on port `8000`. Set `KUMO_RELATIONAL_API_KEY` only when an\n",
    "authenticating gateway protects the endpoint.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f605319a-80ad-45d0-950d-b17fc3e50499",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 138
    },
    "id": "f605319a-80ad-45d0-950d-b17fc3e50499",
    "outputId": "8791220c-bb70-4cff-da6d-0220662a8aee"
   },
   "outputs": [],
   "source": [
    "NIM_URL = os.environ.get(\n",
    "    'KUMO_RELATIONAL_API_ENDPOINT',\n",
    "    'http://localhost:8000',\n",
    ")\n",
    "NIM_API_KEY = os.environ.get('KUMO_RELATIONAL_API_KEY')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "FpYcu6NQYaU5",
   "metadata": {
    "id": "FpYcu6NQYaU5"
   },
   "source": [
    "## Step 3. Connect to the NIM\n",
    "\n",
    "A `RelationalClient` owns the HTTP connection to one NIM endpoint. Keep credentials\n",
    "in environment configuration rather than notebook source.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "r7HX2s1iYaU5",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "r7HX2s1iYaU5",
    "outputId": "25ed5fe9-c257-4471-8a88-61ed4272543c"
   },
   "outputs": [],
   "source": [
    "client = RelationalClient(\n",
    "    url=NIM_URL,\n",
    "    api_key=NIM_API_KEY,\n",
    ")\n",
    "\n",
    "print('Ready:', client.health_ready())\n",
    "print('Capabilities:', client.capabilities('kumo-relational'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0bdd915",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "<strong>Security:</strong> Never hard-code NIM or gateway credentials in a\n",
    "notebook. Use environment variables or your platform's secret manager.\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "RvRgi0DaYaU5",
   "metadata": {
    "id": "RvRgi0DaYaU5"
   },
   "source": [
    "## Step 4. Data Preparation\n",
    "\n",
    "This notebook uses the **AdventureWorks OLTP sample database**, originally\n",
    "published by Microsoft as a SQL Server sample. The four source extracts are:\n",
    "\n",
    "- `Customer.csv`\n",
    "- `Product.csv`\n",
    "- `SalesOrderHeader.csv`\n",
    "- `SalesOrderDetail.csv`\n",
    "\n",
    "Download the four `*.csv` extracts from the AdventureWorks OLTP install scripts in the\n",
    "[Microsoft SQL Server samples repository](https://github.com/microsoft/sql-server-samples/tree/master/samples/databases/adventure-works/oltp-install-script).\n",
    "\n",
    "> Each user is responsible for checking the content of datasets and the\n",
    "> applicable licenses and determining if suitable for the intended use.\n",
    "The repository's samples and templates are provided under the\n",
    "[MIT License](https://github.com/microsoft/sql-server-samples/blob/master/license.txt).\n",
    "\n",
    "Set `ADVENTURE_WORKS_DATA_DIR` to choose where the files are read from,\n",
    "e.g. `export ADVENTURE_WORKS_DATA_DIR=/path/to/adventure-works-cache`. If it is unset,\n",
    "the notebook uses `../rfm/adventure-works` relative to its working directory.\n",
    "The files are tab-delimited and do not include a header row, so each schema is\n",
    "assigned explicitly when the file is loaded.\n",
    "\n",
    "The sales-order header and detail tables remain separate. During preparation,\n",
    "only `CustomerID` and `OrderDate` are copied from each header to its detail\n",
    "rows using `SalesOrderID`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "uGUWXhbsYaU5",
   "metadata": {
    "id": "uGUWXhbsYaU5"
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "DATA_DIR = (\n",
    "    Path(os.environ.get('ADVENTURE_WORKS_DATA_DIR', '../rfm/adventure-works'))\n",
    "    .expanduser()\n",
    "    .resolve()\n",
    ")\n",
    "\n",
    "print('AdventureWorks data:', DATA_DIR)\n",
    "\n",
    "CUSTOMER_COLUMNS = [\n",
    "    'CustomerID',\n",
    "    'PersonID',\n",
    "    'StoreID',\n",
    "    'TerritoryID',\n",
    "    'AccountNumber',\n",
    "    'rowguid',\n",
    "    'ModifiedDate',\n",
    "]\n",
    "\n",
    "PRODUCT_COLUMNS = [\n",
    "    'ProductID',\n",
    "    'Name',\n",
    "    'ProductNumber',\n",
    "    'MakeFlag',\n",
    "    'FinishedGoodsFlag',\n",
    "    'Color',\n",
    "    'SafetyStockLevel',\n",
    "    'ReorderPoint',\n",
    "    'StandardCost',\n",
    "    'ListPrice',\n",
    "    'Size',\n",
    "    'SizeUnitMeasureCode',\n",
    "    'WeightUnitMeasureCode',\n",
    "    'Weight',\n",
    "    'DaysToManufacture',\n",
    "    'ProductLine',\n",
    "    'Class',\n",
    "    'Style',\n",
    "    'ProductSubcategoryID',\n",
    "    'ProductModelID',\n",
    "    'SellStartDate',\n",
    "    'SellEndDate',\n",
    "    'DiscontinuedDate',\n",
    "    'rowguid',\n",
    "    'ModifiedDate',\n",
    "]\n",
    "\n",
    "SALES_HEADER_COLUMNS = [\n",
    "    'SalesOrderID',\n",
    "    'RevisionNumber',\n",
    "    'OrderDate',\n",
    "    'DueDate',\n",
    "    'ShipDate',\n",
    "    'Status',\n",
    "    'OnlineOrderFlag',\n",
    "    'SalesOrderNumber',\n",
    "    'PurchaseOrderNumber',\n",
    "    'AccountNumber',\n",
    "    'CustomerID',\n",
    "    'SalesPersonID',\n",
    "    'TerritoryID',\n",
    "    'BillToAddressID',\n",
    "    'ShipToAddressID',\n",
    "    'ShipMethodID',\n",
    "    'CreditCardID',\n",
    "    'CreditCardApprovalCode',\n",
    "    'CurrencyRateID',\n",
    "    'SubTotal',\n",
    "    'TaxAmt',\n",
    "    'Freight',\n",
    "    'TotalDue',\n",
    "    'Comment',\n",
    "    'rowguid',\n",
    "    'ModifiedDate',\n",
    "]\n",
    "\n",
    "SALES_DETAIL_COLUMNS = [\n",
    "    'SalesOrderID',\n",
    "    'SalesOrderDetailID',\n",
    "    'CarrierTrackingNumber',\n",
    "    'OrderQty',\n",
    "    'ProductID',\n",
    "    'SpecialOfferID',\n",
    "    'UnitPrice',\n",
    "    'UnitPriceDiscount',\n",
    "    'LineTotal',\n",
    "    'rowguid',\n",
    "    'ModifiedDate',\n",
    "]\n",
    "\n",
    "\n",
    "def read_headerless_table(filename, columns):\n",
    "    path = DATA_DIR / filename\n",
    "    if not path.exists():\n",
    "        raise FileNotFoundError(\n",
    "            f'{path} was not found. Download the AdventureWorks CSV files '\n",
    "            'listed above before running this cell.'\n",
    "        )\n",
    "    return pd.read_csv(\n",
    "        path,\n",
    "        sep='\\t',\n",
    "        header=None,\n",
    "        names=columns,\n",
    "        na_values=[''],\n",
    "        low_memory=False,\n",
    "    )\n",
    "\n",
    "\n",
    "customers_df = read_headerless_table('Customer.csv', CUSTOMER_COLUMNS)\n",
    "products_df = read_headerless_table('Product.csv', PRODUCT_COLUMNS)\n",
    "sales_order_headers_df = read_headerless_table(\n",
    "    'SalesOrderHeader.csv', SALES_HEADER_COLUMNS\n",
    ")\n",
    "sales_order_details_df = read_headerless_table(\n",
    "    'SalesOrderDetail.csv', SALES_DETAIL_COLUMNS\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "PFBBD4ZIGXgH",
   "metadata": {
    "id": "PFBBD4ZIGXgH"
   },
   "source": [
    "Parse the timestamp fields, then enrich each sales-order detail with only its parent order's `CustomerID` and `OrderDate`. The header and detail dataframes remain separate.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da29027e-e7df-4784-8801-3ffdb91bb3dd",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 143
    },
    "id": "da29027e-e7df-4784-8801-3ffdb91bb3dd",
    "outputId": "ce251795-afc7-4a6d-cb33-55f1194806ec"
   },
   "outputs": [],
   "source": [
    "timestamp_columns = {\n",
    "    'customers': (customers_df, ['ModifiedDate']),\n",
    "    'products': (\n",
    "        products_df,\n",
    "        ['SellStartDate', 'SellEndDate', 'DiscontinuedDate', 'ModifiedDate'],\n",
    "    ),\n",
    "    'sales_order_headers': (\n",
    "        sales_order_headers_df,\n",
    "        ['OrderDate', 'DueDate', 'ShipDate', 'ModifiedDate'],\n",
    "    ),\n",
    "    'sales_order_details': (sales_order_details_df, ['ModifiedDate']),\n",
    "}\n",
    "\n",
    "for _, (table_df, columns) in timestamp_columns.items():\n",
    "    for column in columns:\n",
    "        table_df[column] = pd.to_datetime(table_df[column], errors='coerce')\n",
    "\n",
    "sales_order_details_df = sales_order_details_df.merge(\n",
    "    sales_order_headers_df[['SalesOrderID', 'CustomerID', 'OrderDate']],\n",
    "    on='SalesOrderID',\n",
    "    how='left',\n",
    "    validate='many_to_one',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "PXUd-a9KYaU6",
   "metadata": {
    "id": "PXUd-a9KYaU6"
   },
   "source": [
    "## Step 5. Build a graph from DataFrames\n",
    "\n",
    "`relational.Graph.from_data()` accepts a dictionary of pandas DataFrames and\n",
    "creates the underlying local table objects automatically. You do not need to\n",
    "construct `LocalTable` objects for the standard workflow.\n",
    "\n",
    "The AdventureWorks relationships are supplied explicitly so the example is\n",
    "reproducible. Kumo Relational still infers column data types and semantic metadata,\n",
    "which you should review before prediction.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "QFJvod3zIOr6",
   "metadata": {
    "id": "QFJvod3zIOr6"
   },
   "source": [
    "Create the graph directly from the four prepared\n",
    "DataFrames:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "_OiiUzNkYaU6",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_OiiUzNkYaU6",
    "outputId": "39c6efad-8488-40dd-9622-6e1a70198ba3",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "graph = relational.Graph.from_data(\n",
    "    {\n",
    "        'customers': customers_df,\n",
    "        'products': products_df,\n",
    "        'sales_order_headers': sales_order_headers_df,\n",
    "        'sales_order_details': sales_order_details_df,\n",
    "    },\n",
    "    edges=[\n",
    "        ('sales_order_headers', 'CustomerID', 'customers'),\n",
    "        ('sales_order_details', 'SalesOrderID', 'sales_order_headers'),\n",
    "        ('sales_order_details', 'ProductID', 'products'),\n",
    "        ('sales_order_details', 'CustomerID', 'customers'),\n",
    "    ],\n",
    "    verbose=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36cf621a",
   "metadata": {},
   "source": [
    "Review the inferred metadata before making any changes.\n",
    "\n",
    "The graph-level view lists every table with its primary key and time columns.\n",
    "The table-level views show each column's data type (`dtype`), semantic type\n",
    "(`stype`), and key or time-column role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0fb8bc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Review the inferred graph and table definitions.\n",
    "graph.print_metadata()\n",
    "\n",
    "for table in graph.tables.values():\n",
    "    table.print_metadata()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bd2e672",
   "metadata": {},
   "source": [
    "If the inferred definition does not match your schema, edit only the\n",
    "metadata that needs correction:\n",
    "\n",
    "- Set a primary key with `graph[\"table_name\"].primary_key = \"column_name\"`.\n",
    "  The SDK automatically assigns the `ID` semantic type to that column.\n",
    "- Set the event-time column with\n",
    "  `graph[\"table_name\"].time_column = \"column_name\"`.\n",
    "- Override one column's semantic type with\n",
    "  `graph[\"table_name\"][\"column_name\"].stype = relational.Stype.<type>`.\n",
    "\n",
    "Use the `relational.Stype` enum so the intended type is explicit. The SDK checks\n",
    "that the semantic type is compatible with the column's data type. Foreign-key\n",
    "columns used in the `edges` argument are already assigned `Stype.ID`.\n",
    "\n",
    "For this known AdventureWorks schema, review identifier-like strings before\n",
    "prediction. Account numbers, product numbers, order numbers, and approval codes\n",
    "identify records rather than categories, so the example marks them `Stype.ID`.\n",
    "Several exceed the NIM's 10,000-value categorical cardinality limit in the full\n",
    "extracts. Product names remain `Stype.text` because their language content is a\n",
    "useful feature, while order quantity is a measured numerical value.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "625b3c1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Correct primary keys that are specific to the known schema.\n",
    "# Assigning a primary key also assigns Stype.ID to that column.\n",
    "primary_keys = {\n",
    "    'customers': 'CustomerID',\n",
    "    'products': 'ProductID',\n",
    "    'sales_order_headers': 'SalesOrderID',\n",
    "    'sales_order_details': 'SalesOrderDetailID',\n",
    "}\n",
    "\n",
    "for table_name, column_name in primary_keys.items():\n",
    "    graph[table_name].primary_key = column_name\n",
    "\n",
    "# Identify the event-time column for tables that contain timestamped events.\n",
    "graph['sales_order_headers'].time_column = 'OrderDate'\n",
    "graph['sales_order_details'].time_column = 'OrderDate'\n",
    "\n",
    "# Correct semantic types that are specific to this known schema. These values\n",
    "# identify business records rather than categories or free-text documents.\n",
    "id_columns = {\n",
    "    'customers': ['AccountNumber'],\n",
    "    'products': ['ProductNumber'],\n",
    "    'sales_order_headers': [\n",
    "        'SalesOrderNumber',\n",
    "        'PurchaseOrderNumber',\n",
    "        'AccountNumber',\n",
    "        'CreditCardApprovalCode',\n",
    "    ],\n",
    "}\n",
    "\n",
    "for table_name, column_names in id_columns.items():\n",
    "    for column_name in column_names:\n",
    "        graph[table_name][column_name].stype = relational.Stype.ID\n",
    "\n",
    "# Product names carry natural-language meaning; order quantity is measured.\n",
    "graph['products']['Name'].stype = relational.Stype.text\n",
    "graph['sales_order_details']['OrderQty'].stype = relational.Stype.numerical"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20078be8",
   "metadata": {},
   "source": [
    "> **Important:** A column can have a datetime data type and the\n",
    "> `timestamp` semantic type without being the table's designated\n",
    "> `time_column`. Set `time_column` only when the column represents the table's\n",
    "> event timeline—the time Kumo Relational should use to order events and construct\n",
    "> temporal context. Administrative timestamps such as `ModifiedDate` do not\n",
    "> necessarily serve that role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbbfcecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# These tables contain datetime columns but do not represent event timelines.\n",
    "graph['customers'].time_column = None\n",
    "graph['products'].time_column = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3af6590b",
   "metadata": {},
   "source": [
    "Review the updated definition to confirm that the corrections were\n",
    "applied:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e721c87f",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.print_metadata()\n",
    "\n",
    "for table in graph.tables.values():\n",
    "    table.print_metadata()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9591ac00",
   "metadata": {},
   "source": [
    "Primary-key assignment automatically gives each key column the `ID`\n",
    "semantic type. Relationships supplied through `edges` do the same for their\n",
    "foreign-key columns, so neither needs a separate semantic-type override.\n",
    "\n",
    "Other available semantic types include `categorical`, `multicategorical`,\n",
    "`numerical`, `text`, `timestamp`, and `unsupported`.\n",
    "\n",
    "Call `graph.validate()` after all metadata changes. An incompatible semantic\n",
    "type, a missing column, or a primary key that is also configured as a time\n",
    "column raises an error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f227599c",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph.validate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "HeISl7332Zgr",
   "metadata": {
    "id": "HeISl7332Zgr"
   },
   "source": [
    "**Quick reference:**\n",
    "\n",
    "1. **`stype` (semantic type)**\n",
    "\n",
    "   A semantic type determines how a column is encoded downstream.\n",
    "\n",
    "   | Type | Explanation | AdventureWorks example |\n",
    "   |---|---|---|\n",
    "   | `\"numerical\"` | A measurable numeric value | `LineTotal`, `OrderQty` |\n",
    "   | `\"categorical\"` | A value from a limited set of categories | `Color`, `ProductLine` |\n",
    "   | `\"multicategorical\"` | Multiple categories stored in one cell | A pipe-delimited tag column |\n",
    "   | `\"ID\"` | A primary key, foreign key, or other identifier | `CustomerID`, `ProductID` |\n",
    "   | `\"text\"` | Natural-language text | A product description |\n",
    "   | `\"timestamp\"` | A point in time | `OrderDate` |\n",
    "   | `\"sequence\"` | Custom embeddings or sequential data | A numeric embedding vector |\n",
    "\n",
    "2. **`primary_key`**\n",
    "\n",
    "   - A primary key uniquely identifies each row.\n",
    "   - Duplicate primary-key values are not supported as distinct entities.\n",
    "   - A primary key can be the destination of a primary-key/foreign-key link.\n",
    "   - A table has at most one primary key, which can consist of one or more columns.\n",
    "\n",
    "3. **`time_column`**\n",
    "\n",
    "   - A time column records when an event occurred.\n",
    "   - Its values must be parseable by `pandas.to_datetime`.\n",
    "   - Each table can have at most one time column.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "T4fEPffRYaU7",
   "metadata": {
    "id": "T4fEPffRYaU7"
   },
   "source": [
    "## Step 6. Inspect the relational graph\n",
    "\n",
    "The graph preserves the four separate AdventureWorks tables:\n",
    "\n",
    "- each sales-order header belongs to one customer;\n",
    "- each sales-order detail belongs to one sales-order header;\n",
    "- each sales-order detail references one product; and\n",
    "- the copied `CustomerID` links each detail row to its customer.\n",
    "\n",
    "Review the registered links before sending prediction requests.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "243f4660-42a4-42b8-b0bb-e90d368ee5af",
   "metadata": {
    "id": "243f4660-42a4-42b8-b0bb-e90d368ee5af"
   },
   "source": [
    "Print the relationships supplied to\n",
    "`Graph.from_data()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4682348b-71ef-47e1-a2b6-e31f434376cd",
   "metadata": {
    "id": "4682348b-71ef-47e1-a2b6-e31f434376cd"
   },
   "outputs": [],
   "source": [
    "graph.print_links()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bpe8yBuaOGIz",
   "metadata": {
    "id": "bpe8yBuaOGIz"
   },
   "source": [
    "Visualize the graph as a Mermaid\n",
    "entity-relationship diagram:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pcwsiX_AYaU7",
   "metadata": {
    "id": "pcwsiX_AYaU7"
   },
   "outputs": [],
   "source": [
    "graph.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9FeVEveUOn9v",
   "metadata": {
    "id": "9FeVEveUOn9v"
   },
   "source": [
    "Validate the completed graph:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "696291cf-a60c-4854-9252-43384c3b3b91",
   "metadata": {
    "id": "696291cf-a60c-4854-9252-43384c3b3b91"
   },
   "outputs": [],
   "source": [
    "graph.validate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "Iegen4I8YaU8",
   "metadata": {
    "id": "Iegen4I8YaU8"
   },
   "source": [
    "## Step 7. Submit prediction requests\n",
    "\n",
    "The graph is a reusable description of the data. For each prediction:\n",
    "\n",
    "1. Write the task in PQL.\n",
    "2. Bind the graph to a model handle with `client.relational(graph)`.\n",
    "3. Call `.predict(query, indices=[...])` on that handle.\n",
    "\n",
    "The SDK returns a pandas DataFrame. Standard applications do not initialize\n",
    "the underlying Kumo Relational driver or model object directly.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hFNiSPODQAjd",
   "metadata": {
    "id": "hFNiSPODQAjd"
   },
   "source": [
    "**Note:** AdventureWorks is sample data, and the queries and results below are intended for demonstration. Benchmark Kumo Relational on data representative of your own use case before drawing performance conclusions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ldqs69MgPgO",
   "metadata": {
    "id": "9ldqs69MgPgO"
   },
   "source": [
    "#### PQL: How you describe a prediction task\n",
    "\n",
    "Learn more in the\n",
    "[PQL documentation](https://docs.nvidia.com/sdgm/reference/pq-reference-overview)\n",
    "and [query-writing guide](https://docs.nvidia.com/sdgm/rfm/writing-predictive-queries).\n",
    "\n",
    "The short reference below covers the entity, target, aggregation,\n",
    "prediction-window, and condition syntax used by the examples in this\n",
    "quickstart.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wiE36SMUNvL9",
   "metadata": {
    "id": "wiE36SMUNvL9"
   },
   "source": [
    "<details>\n",
    "<summary><b>💡 Click here for a short introduction to Predictive Query!</b></summary>\n",
    "\n",
    "**Entities** can be specified in three ways:\n",
    "\n",
    "* `FOR <entity_table>.<entity_primary_key>=1` embeds one entity ID in PQL.\n",
    "* `FOR <entity_table>.<entity_primary_key> IN (1, 2, 3)` embeds several\n",
    "  entity IDs in PQL.\n",
    "* `FOR EACH <entity_table>.<entity_primary_key>` declares the entity column,\n",
    "  while `predict(query, indices=[1, 2, 3])` supplies the entity IDs separately.\n",
    "\n",
    "**Preferred:** For application code, use `FOR EACH` and pass primary-key values\n",
    "through `predict()`'s `indices` argument. This keeps the prediction task separate\n",
    "from the entity batch and makes the same PQL reusable. Despite its name,\n",
    "`indices` contains source primary-key values, not DataFrame row numbers or\n",
    "internal graph indices. If `indices` is provided, it overrides entity IDs\n",
    "embedded in the PQL.\n",
    "\n",
    "**Targets** can be specified in two common ways:\n",
    "\n",
    "* `PREDICT <target_table>.<target_column>` imputes a missing value.\n",
    "* `PREDICT <aggregation>(<target>, <start_offset>, <end_offset>, <time_unit>)` predicts an aggregate in a future window.\n",
    "\n",
    "Supported aggregations include `COUNT`, `SUM`, `AVG`, `MIN`, `MAX`, and `LIST_DISTINCT`. For example, `PREDICT COUNT(sales_order_headers.*, 0, 7, days)` predicts the number of orders in the next seven days.\n",
    "\n",
    "Conditions can turn an aggregate into an event. For example, appending `=0` predicts the probability that no orders occur in the specified window.\n",
    "</details>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f75b92f",
   "metadata": {},
   "source": [
    "### Example 1A: Forecast 30-day product demand\n",
    "\n",
    "Predict how many units of product `ProductID=870` will be ordered in the next 30 days.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcaabebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = (\n",
    "    'PREDICT SUM(sales_order_details.OrderQty, 0, 30, days) '\n",
    "    'FOR EACH products.ProductID'\n",
    ")\n",
    "\n",
    "result = client.relational(graph).predict(\n",
    "    query,\n",
    "    indices=[870],\n",
    "    run_mode='fast',\n",
    ")\n",
    "display(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "580952d0",
   "metadata": {},
   "source": [
    "How to interpret the result:\n",
    "\n",
    "1. `ENTITY`: the product with `ProductID=870`\n",
    "2. `ANCHOR_TIMESTAMP`: the point in time from which the next 30 days are predicted\n",
    "3. `PREDICTION`: the predicted number of units ordered during that window\n",
    "\n",
    "This result can support inventory and demand planning.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24bc7674",
   "metadata": {},
   "source": [
    "### Example 1B: Forecast from a historical anchor time\n",
    "\n",
    "By default, predictions use the maximum timestamp in the temporal graph. Set a historical `anchor_time` to simulate what the prediction would have looked like at that point in time. Kumo Relational uses only information available before the anchor time.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4b78357",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pin the anchor instead of letting it be derived. It is taken from the\n",
    "# data rather than written as a literal: AdventureWorks is re-dated\n",
    "# upstream, so a hardcoded date stops matching the file you just read.\n",
    "anchor_time = sales_order_headers_df['OrderDate'].max() - pd.Timedelta(days=60)\n",
    "\n",
    "result = client.relational(graph).predict(\n",
    "    query,\n",
    "    indices=[870],\n",
    "    run_mode='fast',\n",
    "    anchor_time=anchor_time,\n",
    ")\n",
    "display(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fd753a6",
   "metadata": {},
   "source": [
    "### Example 1C: Forecast for multiple steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48f30c0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = (\n",
    "    'PREDICT SUM(sales_order_details.OrderQty, 0, 30, days) '\n",
    "    'FORECAST 3 TIMEFRAMES FOR EACH products.ProductID'\n",
    ")\n",
    "\n",
    "result = client.relational(graph).predict(\n",
    "    query,\n",
    "    indices=[870],\n",
    "    run_mode='fast',\n",
    ")\n",
    "display(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e44059b9",
   "metadata": {},
   "source": [
    "A multi-step forecast returns one row per timeframe.\n",
    "`FORECAST_STEP` is the one-based sequence number and `PREDICTION` is the\n",
    "forecast value for that period.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63cd46cf",
   "metadata": {},
   "source": [
    "### Example 2: Predict customer inactivity\n",
    "\n",
    "Predict the likelihood that customers `CustomerID=11176` and `CustomerID=11091` will place zero orders in the next 90 days.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b02b047e",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = (\n",
    "    'PREDICT COUNT(sales_order_headers.*, 0, 90, days)=0 '\n",
    "    'FOR EACH customers.CustomerID'\n",
    ")\n",
    "\n",
    "result = client.relational(graph).predict(\n",
    "    query,\n",
    "    indices=[11176, 11091],\n",
    "    run_mode='fast',\n",
    ")\n",
    "display(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afec131e",
   "metadata": {},
   "source": [
    "How to interpret the result:\n",
    "\n",
    "1. `ENTITY`: the customer represented by each `CustomerID`\n",
    "2. `ANCHOR_TIMESTAMP`: the point in time from which the next 90 days are predicted\n",
    "3. `PREDICTION`: whether the zero-order event is predicted to occur\n",
    "4. Probability columns: the estimated likelihood of each outcome\n",
    "\n",
    "This result can support retention or re-engagement workflows.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fbfe1ac",
   "metadata": {},
   "source": [
    "### Example 3: Recommend products\n",
    "\n",
    "Predict the top 10 products that customer `CustomerID=11176` is likely to order in the next 30 days.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99df1717",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = (\n",
    "    'PREDICT LIST_DISTINCT(sales_order_details.ProductID, 0, 30, days) '\n",
    "    'RANK TOP 10 FOR EACH customers.CustomerID'\n",
    ")\n",
    "\n",
    "result = client.relational(graph).predict(\n",
    "    query,\n",
    "    indices=[11176],\n",
    "    run_mode='fast',\n",
    ")\n",
    "display(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c910dd58",
   "metadata": {},
   "source": [
    "How to interpret the result:\n",
    "\n",
    "1. `ENTITY`: the customer with `CustomerID=11176`\n",
    "2. `ANCHOR_TIMESTAMP`: the point in time from which the next 30 days are predicted\n",
    "3. `CLASS`: a recommended `ProductID`\n",
    "4. `SCORE`: a higher score indicates a stronger recommendation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a85305e",
   "metadata": {},
   "source": [
    "### Example 4: Infer a missing product attribute\n",
    "\n",
    "Predict the missing `Color` value for product `ProductID=870`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc3d95cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = 'PREDICT products.Color FOR EACH products.ProductID'\n",
    "\n",
    "result = client.relational(graph).predict(\n",
    "    query,\n",
    "    indices=[870],\n",
    "    run_mode='fast',\n",
    ")\n",
    "display(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "715d1e7d",
   "metadata": {},
   "source": [
    "How to interpret the result:\n",
    "\n",
    "1. `ENTITY`: the product with `ProductID=870`\n",
    "2. `ANCHOR_TIMESTAMP`: the temporal context used for the prediction\n",
    "3. `PREDICTION`: the predicted color category\n",
    "\n",
    "Attribute inference can help enrich incomplete catalog records, but predicted values should be reviewed before being written back to a source system.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d22cca7",
   "metadata": {},
   "source": [
    "## Step 8. Close the client\n",
    "\n",
    "Close the client when the notebook is finished to release its HTTP connection.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39a769de",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
