Remote IO support in cudf-polars

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -74,6 +74,7 @@ dependencies:
 - pytest-benchmark
 - pytest-cases>=3.8.2
 - pytest-cov
+- pytest-httpserver
 - pytest-rerunfailures!=16.0.0
 - pytest-xdist
 - python-confluent-kafka>=2.8.0,<2.9.0a0

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -75,6 +75,7 @@ dependencies:
 - pytest-benchmark
 - pytest-cases>=3.8.2
 - pytest-cov
+- pytest-httpserver
 - pytest-rerunfailures!=16.0.0
 - pytest-xdist
 - python-confluent-kafka>=2.8.0,<2.9.0a0

conda/environments/all_cuda-130_arch-aarch64.yaml

@@ -74,6 +74,7 @@ dependencies:
 - pytest-benchmark
 - pytest-cases>=3.8.2
 - pytest-cov
+- pytest-httpserver
 - pytest-rerunfailures!=16.0.0
 - pytest-xdist
 - python-confluent-kafka>=2.8.0,<2.9.0a0

conda/environments/all_cuda-130_arch-x86_64.yaml

@@ -75,6 +75,7 @@ dependencies:
 - pytest-benchmark
 - pytest-cases>=3.8.2
 - pytest-cov
+- pytest-httpserver
 - pytest-rerunfailures!=16.0.0
 - pytest-xdist
 - python-confluent-kafka>=2.8.0,<2.9.0a0

dependencies.yaml

@@ -891,6 +891,7 @@ dependencies:
       - output_types: [conda, requirements, pyproject]
         packages:
           - rich
+          - pytest-httpserver
   test_python_narwhals:
     common:
       - output_types: [conda, requirements, pyproject]
@@ -1207,6 +1208,7 @@ dependencies:
       - output_types: conda
         packages:
           - cudf-polars==25.10.*,>=0.0.0a0
+          - pytest-httpserver
   depends_on_narwhals:
     common:
       - output_types: [conda, requirements, pyproject]

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -16,6 +16,7 @@
 import itertools
 import json
 import random
+import re
 import time
 from functools import cache
 from pathlib import Path
@@ -268,6 +269,10 @@ class Scan(IR):
     PARQUET_DEFAULT_CHUNK_SIZE: int = 0  # unlimited
     PARQUET_DEFAULT_PASS_LIMIT: int = 16 * 1024**3  # 16GiB
 
+    @staticmethod
+    def _is_url(p: str) -> bool:
+        return re.match(r"^[a-zA-Z][a-zA-Z0-9+.\-]*://", str(p)) is not None
+
     def __init__(
         self,
         schema: Schema,
@@ -325,18 +330,22 @@ def __init__(
             raise NotImplementedError(
                 "Read from cloud storage"
             )  # pragma: no cover; no test yet
-        if any(
-            str(p).startswith("https:/" if POLARS_VERSION_LT_131 else "https://")
-            for p in self.paths
-        ):
+        if (
+            any(str(p).startswith("https:/") for p in self.paths)
+            and POLARS_VERSION_LT_131
+        ):  # pragma: no cover; no test yet
             raise NotImplementedError("Read from https")
         if any(
             str(p).startswith("file:/" if POLARS_VERSION_LT_131 else "file://")
             for p in self.paths
         ):
-            # TODO: removing the file:// may work
             raise NotImplementedError("Read from file URI")
         if self.typ == "csv":
+            if any(self._is_url(p) for p in self.paths):
+                # cannot because we do some file introspection
+                raise NotImplementedError(
+                    "Reading CSV from remote is not yet supported"
+                )
             if self.reader_options["skip_rows_after_header"] != 0:
                 raise NotImplementedError("Skipping rows after header in CSV reader")
             parse_options = self.reader_options["parse_options"]

python/cudf_polars/pyproject.toml

@@ -43,6 +43,7 @@ test = [
     "numpy>=1.23,<3.0a0",
     "pytest",
     "pytest-cov",
+    "pytest-httpserver",
     "pytest-xdist",
     "rich",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.

python/cudf_polars/tests/test_scan.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING
 
 import pytest
+from werkzeug import Response
 
 import polars as pl
 
@@ -18,6 +19,9 @@
 if TYPE_CHECKING:
     from pathlib import Path
 
+    from pytest_httpserver import HTTPServer
+    from werkzeug import Request
+
 
 NO_CHUNK_ENGINE = pl.GPUEngine(raise_on_fail=True, parquet_options={"chunked": False})
 
@@ -487,3 +491,72 @@ def test_scan_from_file_uri(tmp_path: Path) -> None:
     df.write_parquet(path)
     q = pl.scan_parquet(f"file://{path}")
     assert_ir_translation_raises(q, NotImplementedError)
+
+
+@pytest.mark.parametrize("chunked", [False, True])
+def test_scan_parquet_remote(
+    tmp_path: Path, df: pl.DataFrame, httpserver: HTTPServer, *, chunked: bool
+) -> None:
+    path = tmp_path / "foo.parquet"
+    df.write_parquet(path)
+    bytes_ = path.read_bytes()
+    size = len(bytes_)
+
+    def head_handler(_: Request) -> Response:
+        return Response(
+            status=200,
+            headers={
+                "Content-Type": "parquet",
+                "Accept-Ranges": "bytes",
+                "Content-Length": size,
+            },
+        )
+
+    def get_handler(req: Request) -> Response:
+        # parse bytes=200-500 for example (the actual data)
+        start, end = map(int, req.headers["Range"][6:].split("-"))
+        mv = memoryview(bytes_)[start : end + 1]
+        return Response(
+            mv.tobytes(),
+            status=200,
+            headers={
+                "Content-Type": "parquet",
+                "Accept-Ranges": "bytes",
+                "Content-Length": len(mv),
+                "Content-Range": f"bytes {start}-{end - 1}/{size}",
+            },
+        )
+
+    httpserver.expect_request(path, method="HEAD").respond_with_handler(head_handler)
+    httpserver.expect_request(path, method="GET").respond_with_handler(get_handler)
+
+    q = pl.scan_parquet(path)
+
+    assert_gpu_result_equal(
+        q, engine=pl.GPUEngine(raise_on_fail=True, parquet_options={"chunked": chunked})
+    )
+
+
+def test_scan_ndjson_remote(
+    tmp_path: Path, df: pl.LazyFrame, httpserver: HTTPServer
+) -> None:
+    path = tmp_path / "foo.jsonl"
+    df.write_ndjson(path)
+    bytes_ = path.read_bytes()
+    size = len(bytes_)
+
+    def get_handler(_: Request) -> Response:
+        return Response(
+            bytes_,
+            status=200,
+            headers={
+                "Content-Type": "ndjson",
+                "Content-Length": size,
+            },
+        )
+
+    httpserver.expect_request(path, method="HEAD").respond_with_handler(get_handler)
+    httpserver.expect_request(path, method="GET").respond_with_handler(get_handler)
+
+    q = pl.scan_ndjson(path)
+    assert_gpu_result_equal(q)