[BUG] Allow for writes to s3a and s3n paths (#2054)

Closes: #2050

Co-authored-by: Jay Chia <[email protected]@users.noreply.github.com>

daft/filesystem.py

@@ -138,6 +138,10 @@ def _resolve_paths_and_filesystem(
     assert all(isinstance(p, str) for p in paths), paths
     assert len(paths) > 0, paths
 
+    # Sanitize s3a/s3n protocols, which are produced by Hadoop-based systems as a way of denoting which s3
+    # filesystem client to use. However this doesn't matter for Daft, and PyArrow cannot recognize these protocols.
+    paths = [f"s3://{p[6:]}" if p.startswith("s3a://") or p.startswith("s3n://") else p for p in paths]
+
     # Ensure that protocols for all paths are consistent, i.e. that they would map to the
     # same filesystem.
     protocols = {get_protocol_from_path(path) for path in paths}

tests/integration/io/test_write_files_s3_minio.py

@@ -26,15 +26,16 @@ def bucket(minio_io_config):
 
 
 @pytest.mark.integration()
-def test_writing_parquet(minio_io_config, bucket):
+@pytest.mark.parametrize("protocol", ["s3://", "s3a://", "s3n://"])
+def test_writing_parquet(minio_io_config, bucket, protocol):
     data = {
         "foo": [1, 2, 3],
         "bar": ["a", "b", "c"],
     }
     df = daft.from_pydict(data)
     df = df.repartition(2)
     results = df.write_parquet(
-        f"s3://{bucket}/parquet-writes-{uuid.uuid4()}",
+        f"{protocol}{bucket}/parquet-writes-{uuid.uuid4()}",
         partition_cols=["bar"],
         io_config=minio_io_config,
     )