feat: Thumbnail and PDF for new column type (#3193)

* draft: adding pdf cell type

* draft: adding pdf functions for rows processing

* draft: detect modalities

* Rollback

* Fix dependencies

* Adding width and height

* Adding test in assets

* Adding pdf size y bytes

* Add test for get_cell_value_value

* Add more tests

* Fix test parquet and info

* FIx typo

* Update datasets library and add test for modalities

* Add e2e

* Add missing test for first-rows and filter

* Debug e2e

* Fix copy/paste bug

* Fix documents name

* Fix typo

* Address code review suggestions

* Add Pdf spec in OpenApi

Author: AndreaFrancis

docs/source/openapi.json

@@ -392,6 +392,9 @@
           },
           {
             "$ref": "#/components/schemas/VideoFeature"
+          },
+          {
+            "$ref": "#/components/schemas/PdfFeature"
           }
         ]
       },
@@ -577,6 +580,19 @@
           }
         }
       },
+      "PdfFeature": {
+        "type": "object",
+        "required": ["_type"],
+        "properties": {
+          "_type": {
+            "type": "string",
+            "enum": ["Pdf"]
+          },
+          "decode": {
+            "type": "boolean"
+          }
+        }
+      },
       "VideoFeature": {
         "type": "object",
         "required": ["_type"],
@@ -657,6 +673,9 @@
           },
           {
             "$ref": "#/components/schemas/VideoCell"
+          },
+          {
+            "$ref": "#/components/schemas/PdfCell"
           }
         ]
       },
@@ -804,6 +823,22 @@
           ]
         }
       },
+      "PdfCell": {
+        "type": "object",
+        "properties": {
+          "src": {
+            "type": "string",
+            "format": "uri"
+          },
+          "thumbnail":{
+            "$ref": "#/components/schemas/ImageCell"
+          },
+          "size_bytes": {
+            "type": "integer"
+          }
+        },
+        "required": ["src", "thumbnail", "size_bytes"]
+      },
       "VideoCell": {
         "type": "object",
         "properties": {

e2e/tests/conftest.py

@@ -94,3 +94,18 @@ def normal_user_audios_public_dataset() -> Iterator[str]:
         },
     ) as dataset:
         yield dataset
+
+
+@pytest.fixture(scope="session")
+def normal_user_pdfs_public_dataset() -> Iterator[str]:
+    with tmp_dataset(
+        namespace=NORMAL_USER,
+        token=NORMAL_USER_TOKEN,
+        files={
+            "1.pdf": str(Path(__file__).resolve().parent / "data" / "pdfs" / "1.pdf"),
+            "2.pdf": str(Path(__file__).resolve().parent / "data" / "pdfs" / "2.pdf"),
+            "metadata.csv": str(Path(__file__).resolve().parent / "data" / "pdfs" / "metadata.csv"),
+        },
+    ) as dataset:
+        print("Created dataset:", dataset)
+        yield dataset

e2e/tests/data/pdfs/1.pdf

@@ -0,0 +1,3 @@
+file_name,text,has_images
+1.pdf,This is a PDF with an embedded image below,1
+2.pdf,This is a dummy PDF in US Letter size,0

e2e/tests/data/pdfs/2.pdf

@@ -102,3 +102,34 @@ def test_search_audios_endpoint(normal_user_audios_public_dataset: str) -> None:
     # ensure the URL is valid
     response = poll(url, url="")
     assert response.status_code == 200, response
+
+
+def test_search_pdfs_endpoint(normal_user_pdfs_public_dataset: str) -> None:
+    dataset = normal_user_pdfs_public_dataset
+    config, split = get_default_config_split()
+    query = "embedded"
+    rows_response = poll_until_ready_and_assert(
+        relative_url=f"/search?dataset={dataset}&config={config}&split={split}&query={query}",
+        dataset=dataset,
+        should_retry_x_error_codes=["ResponseNotFound"],
+        # ^ I had 404 errors without it. It should return something else at one point.
+    )
+    content = rows_response.json()
+
+    # ensure the URL is signed
+    url = content["rows"][0]["row"]["pdf"]["src"]
+    assert "document.pdf?Expires=" in url, url
+    assert "&Signature=" in url, url
+    assert "&Key-Pair-Id=" in url, url
+    # ensure the URL is valid
+    response = poll(url, url="")
+    assert response.status_code == 200, response
+
+    # ensure the PDF's thumbnail URL is signed
+    thumbnail_url = content["rows"][0]["row"]["pdf"]["thumbnail"]["src"]
+    assert "document.pdf.png?Expires=" in thumbnail_url, thumbnail_url
+    assert "&Signature=" in thumbnail_url, thumbnail_url
+    assert "&Key-Pair-Id=" in thumbnail_url, thumbnail_url
+    # ensure the PDF's thumbnail URL is valid
+    response = poll(thumbnail_url, url="")
+    assert response.status_code == 200, response

e2e/tests/data/pdfs/metadata.csv

@@ -162,3 +162,34 @@ def test_filter_audios_endpoint(normal_user_audios_public_dataset: str) -> None:
     # ensure the URL is valid
     response = poll(url, url="")
     assert response.status_code == 200, response
+
+
+def test_filter_pdfs_endpoint(normal_user_pdfs_public_dataset: str) -> None:
+    dataset = normal_user_pdfs_public_dataset
+    config, split = get_default_config_split()
+    where = "has_images=1"
+    rows_response = poll_until_ready_and_assert(
+        relative_url=f"/filter?dataset={dataset}&config={config}&split={split}&where={where}",
+        dataset=dataset,
+        should_retry_x_error_codes=["ResponseNotFound"],
+        # ^ I had 404 errors without it. It should return something else at one point.
+    )
+    content = rows_response.json()
+
+    # ensure the URL is signed
+    url = content["rows"][0]["row"]["pdf"]["src"]
+    assert "document.pdf?Expires=" in url, url
+    assert "&Signature=" in url, url
+    assert "&Key-Pair-Id=" in url, url
+    # ensure the URL is valid
+    response = poll(url, url="")
+    assert response.status_code == 200, response
+
+    # ensure the PDF's thumbnail URL is signed
+    thumbnail_url = content["rows"][0]["row"]["pdf"]["thumbnail"]["src"]
+    assert "document.pdf.png?Expires=" in thumbnail_url, thumbnail_url
+    assert "&Signature=" in thumbnail_url, thumbnail_url
+    assert "&Key-Pair-Id=" in thumbnail_url, thumbnail_url
+    # ensure the PDF's thumbnail URL is valid
+    response = poll(thumbnail_url, url="")
+    assert response.status_code == 200, response

e2e/tests/test_52_search.py

@@ -81,3 +81,33 @@ def test_rows_audios_endpoint(normal_user_audios_public_dataset: str) -> None:
     # ensure the URL is valid
     response = poll(url, url="")
     assert response.status_code == 200, response
+
+
+def test_rows_pdfs_endpoint(normal_user_pdfs_public_dataset: str) -> None:
+    dataset = normal_user_pdfs_public_dataset
+    config, split = get_default_config_split()
+    rows_response = poll_until_ready_and_assert(
+        relative_url=f"/rows?dataset={dataset}&config={config}&split={split}",
+        dataset=dataset,
+        should_retry_x_error_codes=["ResponseNotFound"],
+        # ^ I had 404 errors without it. It should return something else at one point.
+    )
+    content = rows_response.json()
+
+    # ensure the PDF URL is signed
+    url = content["rows"][0]["row"]["pdf"]["src"]
+    assert "document.pdf?Expires=" in url, url
+    assert "&Signature=" in url, url
+    assert "&Key-Pair-Id=" in url, url
+    # ensure the PDF URL is valid
+    response = poll(url, url="")
+    assert response.status_code == 200, response
+
+    # ensure the PDF's thumbnail URL is signed
+    thumbnail_url = content["rows"][0]["row"]["pdf"]["thumbnail"]["src"]
+    assert "document.pdf.png?Expires=" in thumbnail_url, thumbnail_url
+    assert "&Signature=" in thumbnail_url, thumbnail_url
+    assert "&Key-Pair-Id=" in thumbnail_url, thumbnail_url
+    # ensure the PDF's thumbnail URL is valid
+    response = poll(thumbnail_url, url="")
+    assert response.status_code == 200, response

e2e/tests/test_53_filter.py

@@ -48,3 +48,36 @@ def test_first_rows_audios_endpoint(normal_user_audios_public_dataset: str) -> N
     # ensure the URL is valid
     response = poll(url, url="")
     assert response.status_code == 200, response
+
+
+def test_first_rows_pdfs_endpoint(normal_user_pdfs_public_dataset: str) -> None:
+    dataset = normal_user_pdfs_public_dataset
+    config, split = get_default_config_split()
+    rows_response = poll_until_ready_and_assert(
+        relative_url=f"/first-rows?dataset={dataset}&config={config}&split={split}",
+        dataset=dataset,
+        should_retry_x_error_codes=["ResponseNotFound"],
+        # ^ I had 404 errors without it. It should return something else at one point.
+    )
+    content = rows_response.json()
+
+    # ensure the URL is signed
+    url = content["rows"][0]["row"]["pdf"]["src"]
+    assert isinstance(url, str), url
+    assert "document.pdf?Expires=" in url, url
+    assert "&Signature=" in url, url
+    assert "&Key-Pair-Id=" in url, url
+    # ensure the URL has been signed only once
+    assert url.count("Expires") == 1, url
+    # ensure the URL is valid
+    response = poll(url, url="")
+    assert response.status_code == 200, response
+
+    # ensure the PDF's thumbnail URL is signed
+    thumbnail_url = content["rows"][0]["row"]["pdf"]["thumbnail"]["src"]
+    assert "document.pdf.png?Expires=" in thumbnail_url, thumbnail_url
+    assert "&Signature=" in thumbnail_url, thumbnail_url
+    assert "&Key-Pair-Id=" in thumbnail_url, thumbnail_url
+    # ensure the PDF's thumbnail URL is valid
+    response = poll(thumbnail_url, url="")
+    assert response.status_code == 200, response