From 153f9977455473025b7c6a8a8626a5bce6907200 Mon Sep 17 00:00:00 2001
From: Adam Laiacano <alaiacano@nvidia.com>
Date: Wed, 17 May 2023 20:31:04 -0400
Subject: [PATCH] replace num_rows with head

---
 ...g-Ranking-Models-With-Merlin-Systems.ipynb | 22 +++++++++++++++++--
 ...ving_ranking_models_with_merlin_systems.py |  4 +---
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb b/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb
index ddbb137d9..78121567d 100644
--- a/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb
+++ b/examples/Serving-Ranking-Models-With-Merlin-Systems.ipynb
@@ -27,6 +27,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "f550a0e5",
    "metadata": {},
@@ -70,6 +71,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "1346eec7-48d9-41b3-a3f7-c88230a6f24c",
    "metadata": {},
@@ -90,6 +92,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "222fde5f",
    "metadata": {},
@@ -118,6 +121,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "746fba56-36af-44a3-8beb-34ca026648be",
    "metadata": {},
@@ -150,6 +154,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "143795f2",
    "metadata": {},
@@ -160,6 +165,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "bd5b046f",
    "metadata": {},
@@ -210,6 +216,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "53908458",
    "metadata": {
@@ -248,6 +255,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "8b916afa",
    "metadata": {},
@@ -638,6 +646,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "4edcbd33-f845-4921-b295-46ad366da722",
    "metadata": {},
@@ -667,6 +676,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "96aa55fb",
    "metadata": {},
@@ -735,6 +745,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "0b99146c",
    "metadata": {},
@@ -761,6 +772,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "c75e2eb9",
    "metadata": {},
@@ -800,6 +812,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "358bb59b-796a-4e1e-8e02-b7d3b86c27f9",
    "metadata": {},
@@ -845,6 +858,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "dfbc6f01-5ccf-43a2-ae4f-8451a420fc70",
    "metadata": {},
@@ -1026,12 +1040,13 @@
     "\n",
     "# read in data for request\n",
     "batch = df_lib.read_parquet(\n",
-    "    os.path.join(original_data_path,\"valid\", \"part.0.parquet\"), num_rows=3, columns=workflow.input_schema.column_names\n",
-    ")\n",
+    "    os.path.join(original_data_path,\"valid\", \"part.0.parquet\"), columns=workflow.input_schema.column_names\n",
+    ").head(3)\n",
     "batch"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "9691d28f-80d8-4520-b6e6-b7e260c50380",
    "metadata": {},
@@ -1061,6 +1076,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "50fb442d-d918-4b1a-a825-dfce4f75bc02",
    "metadata": {},
@@ -1081,6 +1097,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "10c00a98-a30e-45ca-9ab6-667621a20c75",
    "metadata": {},
@@ -1111,6 +1128,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "f37053aa-9a01-40ce-ac81-e6edf266f2af",
    "metadata": {},
diff --git a/tests/integration/examples/test_serving_ranking_models_with_merlin_systems.py b/tests/integration/examples/test_serving_ranking_models_with_merlin_systems.py
index d16ab0a0c..2cdbde419 100644
--- a/tests/integration/examples/test_serving_ranking_models_with_merlin_systems.py
+++ b/tests/integration/examples/test_serving_ranking_models_with_merlin_systems.py
@@ -1,7 +1,6 @@
 import os
 
 import pytest
-
 from testbook import testbook
 
 from tests.conftest import REPO_ROOT
@@ -97,9 +96,8 @@ def test_example_04_exporting_ranking_models(tb):
     # read in data for request
     batch = df_lib.read_parquet(
         os.path.join("/tmp/data/", "valid", "part.0.parquet"),
-        num_rows=3,
         columns=workflow.input_schema.column_names,
-    )
+    ).head(3)
     batch = batch.drop(columns="click")
     outputs = tb.ref("output_cols")
     from merlin.dataloader.tf_utils import configure_tensorflow