Copilot commented on code in PR #1625:
URL:
https://github.com/apache/datafusion-python/pull/1625#discussion_r3668676834
##########
python/datafusion/context.py:
##########
@@ -611,6 +613,34 @@ def deregister_object_store(self, schema: str, host: str |
None = None) -> None:
"""
self.ctx.deregister_object_store(schema, host)
+ def _register_object_store_for_path(
+ self, path: str | pathlib.Path, store: Any
+ ) -> None:
+ """Parse a URL path and register the given object store for its scheme
and host.
+
+ This is a convenience helper used by methods like
+ :py:meth:`register_parquet` and :py:meth:`read_parquet` to
+ automatically register an object store when an ``object_store``
+ parameter is provided.
+
+ Args:
+ path: A URL-style path (e.g. ``"s3://bucket/key.parquet"``).
+ store: An object store instance to register.
+
+ Raises:
+ ValueError: If the path does not contain a recognized URL scheme.
+ """
+ parsed = urlparse(str(path))
+ if not parsed.scheme or not parsed.netloc:
+ msg = (
+ f"Cannot determine object store URL from path {path!r}. "
+ "The path must use a URL scheme (e.g. 's3://bucket/key')."
+ )
+ raise ValueError(msg)
+ scheme = f"{parsed.scheme}://"
+ host = parsed.netloc
+ self.register_object_store(scheme, store, host=host)
Review Comment:
`file://` URLs commonly parse with an empty `netloc` (e.g.
`file:///tmp/a.parquet` → `scheme='file'`, `netloc=''`), so this helper will
raise `ValueError` for valid local file URLs. This breaks the intended
`LocalFileSystem` flow (and the new integration tests). Consider allowing empty
`netloc` for `file` (and registering with `host=None`), while continuing to
require `netloc` for schemes like `s3/gs/az/https`.
##########
python/datafusion/context.py:
##########
@@ -611,6 +613,34 @@ def deregister_object_store(self, schema: str, host: str |
None = None) -> None:
"""
self.ctx.deregister_object_store(schema, host)
+ def _register_object_store_for_path(
+ self, path: str | pathlib.Path, store: Any
+ ) -> None:
+ """Parse a URL path and register the given object store for its scheme
and host.
+
+ This is a convenience helper used by methods like
+ :py:meth:`register_parquet` and :py:meth:`read_parquet` to
+ automatically register an object store when an ``object_store``
+ parameter is provided.
+
+ Args:
+ path: A URL-style path (e.g. ``"s3://bucket/key.parquet"``).
+ store: An object store instance to register.
+
+ Raises:
+ ValueError: If the path does not contain a recognized URL scheme.
Review Comment:
The docstring says this raises when the path does not contain a 'recognized
URL scheme', but the implementation doesn't validate schemes and (currently)
also rejects some valid URL forms based on `netloc` being empty. After
adjusting `file://` handling, consider updating the docstring to describe the
actual requirement (e.g., 'must be a URL with a scheme; host/bucket required
for some schemes').
##########
python/datafusion/context.py:
##########
@@ -611,6 +613,34 @@ def deregister_object_store(self, schema: str, host: str |
None = None) -> None:
"""
self.ctx.deregister_object_store(schema, host)
+ def _register_object_store_for_path(
+ self, path: str | pathlib.Path, store: Any
+ ) -> None:
+ """Parse a URL path and register the given object store for its scheme
and host.
Review Comment:
The docstring says this raises when the path does not contain a 'recognized
URL scheme', but the implementation doesn't validate schemes and (currently)
also rejects some valid URL forms based on `netloc` being empty. After
adjusting `file://` handling, consider updating the docstring to describe the
actual requirement (e.g., 'must be a URL with a scheme; host/bucket required
for some schemes').
##########
python/tests/test_object_store_param.py:
##########
@@ -0,0 +1,286 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for the object_store parameter on register/read file methods."""
+
+import contextlib
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pyarrow as pa
+import pytest
+from datafusion import SessionContext
+
+
[email protected]
+def ctx():
+ return SessionContext()
+
+
+class TestRegisterObjectStoreForPath:
+ """Unit tests for _register_object_store_for_path URL parsing logic."""
+
+ def test_parses_s3_url(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ ctx._register_object_store_for_path(
+ "s3://my-bucket/path/to/file.parquet", mock_store
+ )
+ mock_register.assert_called_once_with("s3://", mock_store,
host="my-bucket")
+
+ def test_parses_gs_url(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ ctx._register_object_store_for_path(
+ "gs://my-gcs-bucket/data.parquet", mock_store
+ )
+ mock_register.assert_called_once_with(
+ "gs://", mock_store, host="my-gcs-bucket"
+ )
+
+ def test_parses_az_url(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ ctx._register_object_store_for_path(
+ "az://my-container/data.parquet", mock_store
+ )
+ mock_register.assert_called_once_with(
+ "az://", mock_store, host="my-container"
+ )
+
+ def test_parses_https_url(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ ctx._register_object_store_for_path(
+ "https://my-host.example.com/data.parquet", mock_store
+ )
+ mock_register.assert_called_once_with(
+ "https://", mock_store, host="my-host.example.com"
+ )
+
+ def test_raises_on_local_path(self, ctx):
+ mock_store = MagicMock()
+ with pytest.raises(ValueError, match="Cannot determine object store
URL"):
+ ctx._register_object_store_for_path("/local/path/file.parquet",
mock_store)
+
+ def test_raises_on_relative_path(self, ctx):
+ mock_store = MagicMock()
+ with pytest.raises(ValueError, match="Cannot determine object store
URL"):
+ ctx._register_object_store_for_path("relative/path.parquet",
mock_store)
+
+ def test_raises_on_windows_path(self, ctx):
+ mock_store = MagicMock()
+ with pytest.raises(ValueError, match="Cannot determine object store
URL"):
+ ctx._register_object_store_for_path(
+ "C:\\Users\\data\\file.parquet", mock_store
+ )
+
+ def test_accepts_pathlib_path_raises(self, ctx):
+ """pathlib.Path cannot represent URLs, so this should raise."""
+ mock_store = MagicMock()
+ # pathlib.Path strips the scheme, so this becomes a local path
+ with pytest.raises(ValueError, match="Cannot determine object store
URL"):
+ ctx._register_object_store_for_path(Path("/local/file.parquet"),
mock_store)
+
+
+class TestRegisterParquetObjectStore:
+ """Tests for register_parquet with object_store parameter."""
+
+ def test_object_store_none_does_not_register(self, ctx):
+ """When object_store is None, register_object_store is not called."""
+ with patch.object(ctx, "register_object_store") as mock_register:
+ # This will fail at the Rust level (file doesn't exist), but
+ # we're testing that register_object_store is NOT called
+ with contextlib.suppress(Exception):
+ ctx.register_parquet("t", "s3://bucket/file.parquet")
+ mock_register.assert_not_called()
+
+ def test_object_store_triggers_registration(self, ctx):
+ """When object_store is provided, register_object_store is called."""
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.register_parquet(
+ "t",
+ "s3://my-bucket/file.parquet",
+ object_store=mock_store,
+ )
+ mock_register.assert_called_once_with("s3://", mock_store,
host="my-bucket")
+
+ def test_object_store_invalid_path_raises(self, ctx):
+ """Providing object_store with a local path raises ValueError."""
+ mock_store = MagicMock()
+ with pytest.raises(ValueError, match="Cannot determine object store
URL"):
+ ctx.register_parquet("t", "/local/file.parquet",
object_store=mock_store)
+
+
+class TestReadParquetObjectStore:
+ """Tests for read_parquet with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.read_parquet("s3://my-bucket/file.parquet",
object_store=mock_store)
+ mock_register.assert_called_once_with("s3://", mock_store,
host="my-bucket")
+
+
+class TestRegisterCsvObjectStore:
+ """Tests for register_csv with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.register_csv(
+ "t", "s3://my-bucket/data.csv", object_store=mock_store
+ )
+ mock_register.assert_called_once_with("s3://", mock_store,
host="my-bucket")
+
+ def test_object_store_with_list_path(self, ctx):
+ """For list paths, the first entry is used for URL parsing."""
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.register_csv(
+ "t",
+ ["s3://my-bucket/a.csv", "s3://my-bucket/b.csv"],
+ object_store=mock_store,
+ )
+ mock_register.assert_called_once_with("s3://", mock_store,
host="my-bucket")
+
+
+class TestReadCsvObjectStore:
+ """Tests for read_csv with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.read_csv("gs://bucket/data.csv", object_store=mock_store)
+ mock_register.assert_called_once_with("gs://", mock_store,
host="bucket")
+
+
+class TestRegisterJsonObjectStore:
+ """Tests for register_json with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.register_json("t", "s3://bucket/data.json",
object_store=mock_store)
+ mock_register.assert_called_once_with("s3://", mock_store,
host="bucket")
+
+
+class TestReadJsonObjectStore:
+ """Tests for read_json with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.read_json("s3://bucket/data.json", object_store=mock_store)
+ mock_register.assert_called_once_with("s3://", mock_store,
host="bucket")
+
+
+class TestRegisterAvroObjectStore:
+ """Tests for register_avro with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.register_avro("t", "s3://bucket/data.avro",
object_store=mock_store)
+ mock_register.assert_called_once_with("s3://", mock_store,
host="bucket")
+
+
+class TestReadAvroObjectStore:
+ """Tests for read_avro with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.read_avro("s3://bucket/data.avro", object_store=mock_store)
+ mock_register.assert_called_once_with("s3://", mock_store,
host="bucket")
+
+
+class TestRegisterArrowObjectStore:
+ """Tests for register_arrow with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.register_arrow(
+ "t", "s3://bucket/data.arrow", object_store=mock_store
+ )
+ mock_register.assert_called_once_with("s3://", mock_store,
host="bucket")
+
+
+class TestReadArrowObjectStore:
+ """Tests for read_arrow with object_store parameter."""
+
+ def test_object_store_triggers_registration(self, ctx):
+ mock_store = MagicMock()
+ with patch.object(ctx, "register_object_store") as mock_register:
+ with contextlib.suppress(Exception):
+ ctx.read_arrow("s3://bucket/data.arrow",
object_store=mock_store)
+ mock_register.assert_called_once_with("s3://", mock_store,
host="bucket")
+
+
+class TestEndToEndWithLocalFileSystem:
+ """Integration test using LocalFileSystem object store with
register_parquet."""
+
+ def test_register_parquet_with_local_object_store(self, ctx, tmp_path):
+ """Verify the full flow works with a real object store and local
file."""
+ import pyarrow.parquet as pq
+ from datafusion.object_store import LocalFileSystem
+
+ # Write a test parquet file
+ table = pa.table({"x": [1, 2, 3], "y": ["a", "b", "c"]})
+ parquet_path = tmp_path / "test.parquet"
+ pq.write_table(table, str(parquet_path))
+
+ # Use file:// URL with LocalFileSystem object store
+ store = LocalFileSystem()
+ file_url = f"file://{tmp_path}/test.parquet"
+
+ ctx.register_parquet("test_tbl", file_url, object_store=store)
Review Comment:
Building `file://` URLs via string interpolation is not portable (notably on
Windows, where paths contain backslashes and drive letters) and can produce
invalid URIs. Prefer generating the URL with `Path(...).as_uri()` (e.g., from
`parquet_path`) so the test reliably produces a valid `file://...` URI across
platforms.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]