Repository: ibis-project/ibis

Issue:
does duckdb support for s3:// use fsspec or duckdb's native mechanism?
### Is your feature request related to a problem?

Queries generated by ibis against remote S3 resources using duckdb are taking about twice as long to run for me as when I call `duckdb.sql` with a SQL string directly (using duckdb's `httpfs` extension, not the fsspec support).  I was just wondering if the cause for this might be that ibis is passing the I/O through an additional python layer? (I notice considerably lower bandwidth load when the query is running in ibis too, data transfer rates of around 200 Mb/s instead of around 500 Mb/s).  




### Describe the solution you'd like

If ibis is currently using fsspec to handle the queries, would you consider using duckdb's native httpfs extension instead? I think it can provide considerably better performance in some cases.

### What version of ibis are you running?

4.1.0

### What backend(s) are you using, if any?

duckdb

### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
We're using `pyarrow.dataset` to handle S3 authentication, because DuckDB only supports s3 auth via a single, home-rolled, method of specifying credentials.
I'm surprised it's that much slower, though.


It might be a little horrible, but we could try to use the duckdb `httpfs` builtin by default, then fall back to `pyarrow` if it fails (to try with credentials).  

xref https://github.com/duckdb/duckdb/issues/4021

@cboettig -- are these S3 resources public?  Or do they require authentication?
Thanks! yup, these are public resources, example benchmark below.  I may also be completely wrong about the cause for the speed differences here, it may just be due to Ibis generating less efficient SQL?  This should be a reproducible example on a moderate-sized public dataset, takes about 7-8min for 'pure duckdb' but abour 20min for ibis version on my machine:

## duckdb / sql

```
import duckdb
con = duckdb.connect()
con.execute("INSTALL httpfs")
con.execute("LOAD httpfs")

query = f'''
SELECT class, longitude, latitude, COALESCE(n, 0.0) AS n
FROM (
  SELECT class, longitude, latitude, COUNT(*) AS n
  FROM (
    SELECT DISTINCT
      class,
      genus,
      ROUND(decimallongitude,0) AS longitude,
      ROUND(decimallatitude, 0) AS latitude
    FROM read_parquet("s3://gbif-open-data-us-east-1/occurrence/2023-02-01/occurrence.parquet/*")
    WHERE (phylum = 'Chordata')
  )
  GROUP BY class, longitude, latitude
)
'''

df = con.execute(query).df()

```


## Ibis

```
con = ibis.duckdb.connect()
bucket = "gbif-open-data-us-east-1"
path = "occurrence/2023-02-01/occurrence.parquet"
gbif = con.register(f"s3://{bucket}/{path}")

tbl = (gbif
       .filter(gbif.phylum == "Chordata")
       .mutate(latitude = gbif.decimallatitude.round(0),
               longitude = gbif.decimallongitude.round(0))
       .select(["class", "genus", "longitude", "latitude"])
       .distinct()
       .group_by(["class","longitude", "latitude"])
       .size()
      )
#ibis.show_sql(tbl) # optionally peek at the SQL query
df = tbl.to_pandas()
```

Ok -- I'm guessing `dataset` is creating much more overhead than we expected, but I need to get better timings first.  

If you use `https` links to the buckets, then Ibis uses the `duckdb` `httpfs` plugin and the timings are more in-line with what I expect (a constant time difference of ibis vs. duckdb for compilation overhead)

(I also trimmed this down to just 3 of those parquet files to speed things up but the performance difference is consistent when I ramp up to 10)

```python
[ins] In [4]: con = ibis.duckdb.connect()
         ...: bucket = "gbif-open-data-af-south-1"
         ...: path = "occurrence/2023-02-01/occurrence.parquet"
         ...: gbif = con.read_parquet(
         ...:     [f"https://{bucket}.s3.af-south-1.amazonaws.com/{path}/00000{i}" for i in range(3)]
         ...: )
         ...: 
         ...: tbl = (
         ...:     gbif.filter(gbif.phylum == "Chordata")
         ...:     .mutate(
         ...:         latitude=gbif.decimallatitude.round(0), longitude=gbif.decimallongitude.round(0)
         ...:     )
         ...:     .select(["class", "genus", "longitude", "latitude"])
         ...:     .distinct()
         ...:     .group_by(["class", "longitude", "latitude"])
         ...:     .size()
         ...: )

[ins] In [5]: %time df = tbl.execute()
100% ▕████████████████████████████████████████████████████████████▏ 
CPU times: user 486 ms, sys: 119 ms, total: 605 ms
Wall time: 19.2 s

[ins] In [6]: %time df = tbl.execute()
100% ▕████████████████████████████████████████████████████████████▏ 
CPU times: user 1.42 s, sys: 589 ms, total: 2.01 s
Wall time: 20.3 s
```

```python
[ins] In [4]: import duckdb
         ...: 
         ...: con = duckdb.connect()
         ...: con.execute("INSTALL httpfs")
         ...: con.execute("LOAD httpfs")
         ...: 
         ...: files = [
         ...:     f"https://gbif-open-data-af-south-1.s3.af-south-1.amazonaws.com/occurrence/2023-02-01/oc
         ...: currence.parquet/00000{i}"
         ...:     for i in range(3)
         ...: ]
         ...: 
         ...: query = f"""
         ...: SELECT class, longitude, latitude, COALESCE(n, 0.0) AS n
         ...: FROM (
         ...:   SELECT class, longitude, latitude, COUNT(*) AS n
         ...:   FROM (
         ...:     SELECT DISTINCT
         ...:       class,
         ...:       genus,
         ...:       ROUND(decimallongitude,0) AS longitude,
         ...:       ROUND(decimallatitude, 0) AS latitude
         ...:     FROM read_parquet({files})
         ...:     WHERE (phylum = 'Chordata')
         ...:   )
         ...:   GROUP BY class, longitude, latitude
         ...: )
         ...: """

[ins] In [5]: %time df = con.execute(query).df()
100% ▕████████████████████████████████████████████████████████████▏ 
CPU times: user 9.09 s, sys: 4.08 s, total: 13.2 s
Wall time: 17.7 s

[ins] In [6]: %time df = con.execute(query).df()
100% ▕████████████████████████████████████████████████████████████▏ 
CPU times: user 2.72 s, sys: 1.4 s, total: 4.12 s
Wall time: 18.9 s
```

Resolving patch:
diff --git a/ibis/backends/duckdb/__init__.py b/ibis/backends/duckdb/__init__.py
--- a/ibis/backends/duckdb/__init__.py
+++ b/ibis/backends/duckdb/__init__.py
@@ -9,6 +9,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, MutableMapping
 
+import duckdb
 import sqlalchemy as sa
 import toolz
 
@@ -22,7 +23,6 @@
 from ibis.backends.duckdb.datatypes import parse
 
 if TYPE_CHECKING:
-    import duckdb
     import pandas as pd
     import pyarrow as pa
 
@@ -340,47 +340,65 @@ def read_parquet(
         ir.Table
             The just-registered table
         """
-        import sqlalchemy_views as sav
-
         source_list = normalize_filenames(source_list)
 
-        if any(source.startswith("s3://") for source in source_list):
-            if len(source_list) > 1:
-                raise ValueError("only single s3 paths are supported")
+        table_name = table_name or f"ibis_read_parquet_{next(pa_n)}"
 
-            import pyarrow.dataset as ds
+        # Default to using the native duckdb parquet reader
+        # If that fails because of auth issues, fall back to ingesting via
+        # pyarrow dataset
+        try:
+            self._read_parquet_duckdb_native(source_list, table_name, **kwargs)
+        except sa.exc.OperationalError as e:
+            if isinstance(e.orig, duckdb.IOException):
+                self._read_parquet_pyarrow_dataset(source_list, table_name, **kwargs)
+            else:
+                raise e
 
-            dataset = ds.dataset(source_list[0])
-            table_name = table_name or f"ibis_read_parquet_{next(pa_n)}"
-            self._load_extensions(["httpfs"])
-            # We don't create a view since DuckDB special cases Arrow Datasets
-            # so if we also create a view we end up with both a "lazy table"
-            # and a view with the same name
-            with self.begin() as con:
-                # DuckDB normally auto-detects Arrow Datasets that are defined
-                # in local variables but the `dataset` variable won't be local
-                # by the time we execute against this so we register it
-                # explicitly.
-                con.connection.register(table_name, dataset)
-        else:
-            if any(
-                source.startswith(("http://", "https://")) for source in source_list
-            ):
-                self._load_extensions(["httpfs"])
+        return self.table(table_name)
 
-            if table_name is None:
-                table_name = f"ibis_read_parquet_{next(pa_n)}"
+    def _read_parquet_duckdb_native(
+        self,
+        source_list: str | Iterable[str],
+        table_name: str,
+        **kwargs: Any,
+    ) -> None:
+        import sqlalchemy_views as sav
 
-            source = sa.select(sa.literal_column("*")).select_from(
-                sa.func.read_parquet(
-                    sa.func.list_value(*source_list), _format_kwargs(kwargs)
-                )
+        if any(
+            source.startswith(("http://", "https://", "s3://"))
+            for source in source_list
+        ):
+            self._load_extensions(["httpfs"])
+
+        source = sa.select(sa.literal_column("*")).select_from(
+            sa.func.read_parquet(
+                sa.func.list_value(*source_list), _format_kwargs(kwargs)
             )
-            view = sav.CreateView(sa.table(table_name), source, or_replace=True)
-            with self.begin() as con:
-                con.execute(view)
+        )
+        view = sav.CreateView(sa.table(table_name), source, or_replace=True)
+        with self.begin() as con:
+            con.execute(view)
 
-        return self.table(table_name)
+    def _read_parquet_pyarrow_dataset(
+        self,
+        source_list: str | Iterable[str],
+        table_name: str,
+        **kwargs: Any,
+    ) -> None:
+        import pyarrow.dataset as ds
+
+        dataset = ds.dataset(list(map(ds.dataset, source_list)), **kwargs)
+        self._load_extensions(["httpfs"])
+        # We don't create a view since DuckDB special cases Arrow Datasets
+        # so if we also create a view we end up with both a "lazy table"
+        # and a view with the same name
+        with self.begin() as con:
+            # DuckDB normally auto-detects Arrow Datasets that are defined
+            # in local variables but the `dataset` variable won't be local
+            # by the time we execute against this so we register it
+            # explicitly.
+            con.connection.register(table_name, dataset)
 
     def read_in_memory(
         self, dataframe: pd.DataFrame | pa.Table, table_name: str | None = None


Test patch:
diff --git a/ibis/backends/duckdb/tests/test_register.py b/ibis/backends/duckdb/tests/test_register.py
--- a/ibis/backends/duckdb/tests/test_register.py
+++ b/ibis/backends/duckdb/tests/test_register.py
@@ -192,3 +192,20 @@ def test_set_temp_dir(tmp_path):
     path = tmp_path / "foo" / "bar"
     ibis.duckdb.connect(temp_directory=path)
     assert path.exists()
+
+
+def test_s3_403_fallback(httpserver, monkeypatch):
+    con = ibis.duckdb.connect()
+
+    # monkeypatch to avoid downloading extensions in tests
+    monkeypatch.setattr(con, "_load_extensions", lambda x: True)
+
+    # Throw a 403 to trigger fallback to pyarrow.dataset
+    httpserver.expect_request("/myfile").respond_with_data(
+        "Forbidden", status=403, content_type="text/plain"
+    )
+
+    # Since the URI is nonsense to pyarrow, expect an error, but raises from
+    # pyarrow, which indicates the fallback worked
+    with pytest.raises(pa.lib.ArrowInvalid):
+        con.read_parquet(httpserver.url_for("/myfile"))
diff --git a/ibis/backends/tests/test_register.py b/ibis/backends/tests/test_register.py
--- a/ibis/backends/tests/test_register.py
+++ b/ibis/backends/tests/test_register.py
@@ -369,14 +369,15 @@ def test_csv_reregister_schema(con, tmp_path):
         "trino",
     ]
 )
-def test_register_garbage(con):
+def test_register_garbage(con, monkeypatch):
+    # monkeypatch to avoid downloading extensions in tests
+    monkeypatch.setattr(con, "_load_extensions", lambda x: True)
+
     sa = pytest.importorskip("sqlalchemy")
     with pytest.raises(
         sa.exc.OperationalError, match="No files found that match the pattern"
     ):
         con.read_csv("garbage_notafile")
 
-    with pytest.raises(
-        sa.exc.OperationalError, match="No files found that match the pattern"
-    ):
+    with pytest.raises(FileNotFoundError):
         con.read_parquet("garbage_notafile")