From b225b3a2d6774af3a49ae9f34aa5af315b510317 Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Wed, 8 Jul 2026 12:41:35 -0700 Subject: [PATCH 1/3] Add object_store parameter to register/read file methods Add an optional object_store parameter to register_parquet, read_parquet, register_csv, read_csv, register_json, read_json, register_avro, read_avro, register_arrow, and read_arrow methods. When provided, the store is automatically registered for the URL scheme and host parsed from the path, removing the need to manually call register_object_store separately. This enables thread-safe cloud credential handling without os.environ mutation. Includes unit tests (mock-based) for URL parsing logic and integration tests using LocalFileSystem for end-to-end validation. Closes #1624 --- python/datafusion/context.py | 156 ++++++++++++- python/tests/test_object_store_param.py | 286 ++++++++++++++++++++++++ 2 files changed, 439 insertions(+), 3 deletions(-) create mode 100644 python/tests/test_object_store_param.py diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 0bfc59bfe..2d658aa3a 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -54,6 +54,8 @@ from typing_extensions import deprecated # Python 3.12 +from urllib.parse import urlparse + import pyarrow as pa from datafusion.catalog import ( @@ -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) + def register_listing_table( self, name: str, @@ -1028,6 +1058,7 @@ def register_parquet( skip_metadata: bool = True, schema: pa.Schema | None = None, file_sort_order: Sequence[Sequence[SortKey]] | None = None, + object_store: Any | None = None, ) -> None: """Register a Parquet file as a table. @@ -1049,7 +1080,41 @@ def register_parquet( file_sort_order: Sort order for the file. Each sort key can be specified as a column name (``str``), an expression (``Expr``), or a ``SortExpr``. - """ + object_store: A pre-configured object store instance (e.g. + :py:class:`~datafusion.object_store.AmazonS3`, + :py:class:`~datafusion.object_store.GoogleCloud`, + :py:class:`~datafusion.object_store.MicrosoftAzure`) to use + for accessing the file. When provided, the store is + automatically registered for the URL scheme and host parsed + from ``path``, removing the need to call + :py:meth:`register_object_store` separately. This is + especially useful in multi-threaded environments where + setting credentials via ``os.environ`` is not thread-safe. + + Examples: + Register a local Parquet file: + + >>> import datafusion + >>> ctx = datafusion.SessionContext() + >>> ctx.register_parquet("my_table", "data.parquet") + + Register from S3 with inline credentials (thread-safe): + + >>> from datafusion.object_store import AmazonS3 # doctest: +SKIP + >>> store = AmazonS3( + ... bucket_name="my-bucket", + ... region="us-east-1", + ... access_key_id="...", + ... secret_access_key="...", + ... ) # doctest: +SKIP + >>> ctx.register_parquet( + ... "my_table", + ... "s3://my-bucket/data.parquet", + ... object_store=store, + ... ) # doctest: +SKIP + """ + if object_store is not None: + self._register_object_store_for_path(path, object_store) if table_partition_cols is None: table_partition_cols = [] table_partition_cols = _convert_table_partition_cols(table_partition_cols) @@ -1075,6 +1140,7 @@ def register_csv( file_extension: str = ".csv", file_compression_type: str | None = None, options: CsvReadOptions | None = None, + object_store: Any | None = None, ) -> None: """Register a CSV file as a table. @@ -1096,7 +1162,14 @@ def register_csv( file_compression_type: File compression type. options: Set advanced options for CSV reading. This cannot be combined with any of the other options in this method. - """ + object_store: A pre-configured object store instance to use for + accessing the file. When provided, the store is automatically + registered for the URL scheme and host parsed from ``path``. + """ + if object_store is not None: + # For list paths, register from the first entry + register_path = path[0] if isinstance(path, list) else path + self._register_object_store_for_path(register_path, object_store) if options is not None and ( schema is not None or not has_header @@ -1143,6 +1216,7 @@ def register_json( file_extension: str = ".json", table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_compression_type: str | None = None, + object_store: Any | None = None, ) -> None: """Register a JSON file as a table. @@ -1159,7 +1233,12 @@ def register_json( selected for data input. table_partition_cols: Partition columns. file_compression_type: File compression type. + object_store: A pre-configured object store instance to use for + accessing the file. When provided, the store is automatically + registered for the URL scheme and host parsed from ``path``. """ + if object_store is not None: + self._register_object_store_for_path(path, object_store) if table_partition_cols is None: table_partition_cols = [] table_partition_cols = _convert_table_partition_cols(table_partition_cols) @@ -1180,6 +1259,7 @@ def register_avro( schema: pa.Schema | None = None, file_extension: str = ".avro", table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, + object_store: Any | None = None, ) -> None: """Register an Avro file as a table. @@ -1192,7 +1272,12 @@ def register_avro( schema: The data source schema. file_extension: File extension to select. table_partition_cols: Partition columns. + object_store: A pre-configured object store instance to use for + accessing the file. When provided, the store is automatically + registered for the URL scheme and host parsed from ``path``. """ + if object_store is not None: + self._register_object_store_for_path(path, object_store) if table_partition_cols is None: table_partition_cols = [] table_partition_cols = _convert_table_partition_cols(table_partition_cols) @@ -1205,6 +1290,7 @@ def register_arrow( schema: pa.Schema | None = None, file_extension: str = ".arrow", table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, + object_store: Any | None = None, ) -> None: """Register an Arrow IPC file as a table. @@ -1217,6 +1303,9 @@ def register_arrow( schema: The data source schema. file_extension: File extension to select. table_partition_cols: Partition columns. + object_store: A pre-configured object store instance to use for + accessing the file. When provided, the store is automatically + registered for the URL scheme and host parsed from ``path``. Examples: >>> import tempfile, os @@ -1271,6 +1360,8 @@ def register_arrow( 30 ] """ + if object_store is not None: + self._register_object_store_for_path(path, object_store) if table_partition_cols is None: table_partition_cols = [] table_partition_cols = _convert_table_partition_cols(table_partition_cols) @@ -1690,6 +1781,7 @@ def read_json( file_extension: str = ".json", table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_compression_type: str | None = None, + object_store: Any | None = None, ) -> DataFrame: """Read a line-delimited JSON data source. @@ -1702,10 +1794,15 @@ def read_json( selected for data input. table_partition_cols: Partition columns. file_compression_type: File compression type. + object_store: A pre-configured object store instance to use for + accessing the file. When provided, the store is automatically + registered for the URL scheme and host parsed from ``path``. Returns: DataFrame representation of the read JSON files. """ + if object_store is not None: + self._register_object_store_for_path(path, object_store) if table_partition_cols is None: table_partition_cols = [] table_partition_cols = _convert_table_partition_cols(table_partition_cols) @@ -1731,6 +1828,7 @@ def read_csv( table_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_compression_type: str | None = None, options: CsvReadOptions | None = None, + object_store: Any | None = None, ) -> DataFrame: """Read a CSV data source. @@ -1750,10 +1848,16 @@ def read_csv( file_compression_type: File compression type. options: Set advanced options for CSV reading. This cannot be combined with any of the other options in this method. + object_store: A pre-configured object store instance to use for + accessing the file. When provided, the store is automatically + registered for the URL scheme and host parsed from ``path``. Returns: DataFrame representation of the read CSV files """ + if object_store is not None: + register_path = path[0] if isinstance(path, list) else path + self._register_object_store_for_path(register_path, object_store) if options is not None and ( schema is not None or not has_header @@ -1803,6 +1907,7 @@ def read_parquet( skip_metadata: bool = True, schema: pa.Schema | None = None, file_sort_order: Sequence[Sequence[SortKey]] | None = None, + object_store: Any | None = None, ) -> DataFrame: """Read a Parquet source into a :py:class:`~datafusion.dataframe.Dataframe`. @@ -1822,10 +1927,43 @@ def read_parquet( file_sort_order: Sort order for the file. Each sort key can be specified as a column name (``str``), an expression (``Expr``), or a ``SortExpr``. + object_store: A pre-configured object store instance (e.g. + :py:class:`~datafusion.object_store.AmazonS3`, + :py:class:`~datafusion.object_store.GoogleCloud`, + :py:class:`~datafusion.object_store.MicrosoftAzure`) to use + for accessing the file. When provided, the store is + automatically registered for the URL scheme and host parsed + from ``path``, removing the need to call + :py:meth:`register_object_store` separately. This is + especially useful in multi-threaded environments where + setting credentials via ``os.environ`` is not thread-safe. Returns: DataFrame representation of the read Parquet files - """ + + Examples: + Read a local Parquet file: + + >>> import datafusion + >>> ctx = datafusion.SessionContext() + >>> df = ctx.read_parquet("data.parquet") # doctest: +SKIP + + Read from S3 with inline credentials (thread-safe): + + >>> from datafusion.object_store import AmazonS3 # doctest: +SKIP + >>> store = AmazonS3( + ... bucket_name="my-bucket", + ... region="us-east-1", + ... access_key_id="...", + ... secret_access_key="...", + ... ) # doctest: +SKIP + >>> df = ctx.read_parquet( + ... "s3://my-bucket/data.parquet", + ... object_store=store, + ... ) # doctest: +SKIP + """ + if object_store is not None: + self._register_object_store_for_path(path, object_store) if table_partition_cols is None: table_partition_cols = [] table_partition_cols = _convert_table_partition_cols(table_partition_cols) @@ -1848,6 +1986,7 @@ def read_avro( schema: pa.Schema | None = None, file_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, file_extension: str = ".avro", + object_store: Any | None = None, ) -> DataFrame: """Create a :py:class:`DataFrame` for reading Avro data source. @@ -1856,10 +1995,15 @@ def read_avro( schema: The data source schema. file_partition_cols: Partition columns. file_extension: File extension to select. + object_store: A pre-configured object store instance to use for + accessing the file. When provided, the store is automatically + registered for the URL scheme and host parsed from ``path``. Returns: DataFrame representation of the read Avro file """ + if object_store is not None: + self._register_object_store_for_path(path, object_store) if file_partition_cols is None: file_partition_cols = [] file_partition_cols = _convert_table_partition_cols(file_partition_cols) @@ -1873,6 +2017,7 @@ def read_arrow( schema: pa.Schema | None = None, file_extension: str = ".arrow", file_partition_cols: list[tuple[str, str | pa.DataType]] | None = None, + object_store: Any | None = None, ) -> DataFrame: """Create a :py:class:`DataFrame` for reading an Arrow IPC data source. @@ -1881,6 +2026,9 @@ def read_arrow( schema: The data source schema. file_extension: File extension to select. file_partition_cols: Partition columns. + object_store: A pre-configured object store instance to use for + accessing the file. When provided, the store is automatically + registered for the URL scheme and host parsed from ``path``. Returns: DataFrame representation of the read Arrow IPC file. @@ -1932,6 +2080,8 @@ def read_arrow( 3 ] """ + if object_store is not None: + self._register_object_store_for_path(path, object_store) if file_partition_cols is None: file_partition_cols = [] file_partition_cols = _convert_table_partition_cols(file_partition_cols) diff --git a/python/tests/test_object_store_param.py b/python/tests/test_object_store_param.py new file mode 100644 index 000000000..b57d9f287 --- /dev/null +++ b/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 + + +@pytest.fixture +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) + result = ctx.sql("SELECT * FROM test_tbl").collect() + + assert len(result) == 1 + assert result[0].num_rows == 3 + assert result[0].column("x").to_pylist() == [1, 2, 3] + + def test_read_parquet_with_local_object_store(self, ctx, tmp_path): + """Verify read_parquet works with object_store parameter.""" + import pyarrow.parquet as pq + from datafusion.object_store import LocalFileSystem + + table = pa.table({"val": [10, 20, 30]}) + parquet_path = tmp_path / "read_test.parquet" + pq.write_table(table, str(parquet_path)) + + store = LocalFileSystem() + file_url = f"file://{tmp_path}/read_test.parquet" + + df = ctx.read_parquet(file_url, object_store=store) + result = df.collect() + + assert len(result) == 1 + assert result[0].column("val").to_pylist() == [10, 20, 30] From c3a97f6ac2438fc5f7082afb5b90d3032f335f8a Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Wed, 29 Jul 2026 16:57:26 -0700 Subject: [PATCH 2/3] Fix file:// URL handling in _register_object_store_for_path - Allow empty netloc for file:// scheme (e.g. file:///tmp/data.parquet) - Require netloc (host/bucket) only for non-file schemes (s3, gs, az, https) - Add doctest +SKIP to register_parquet example referencing nonexistent file - Use Path.as_uri() in end-to-end tests for portable file:// URL construction - Add tests for file:// acceptance and scheme-without-host rejection --- python/datafusion/context.py | 21 ++++++++++++++++----- python/tests/test_object_store_param.py | 21 ++++++++++++++++++--- 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 2d658aa3a..94b2bb1c6 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -624,21 +624,32 @@ def _register_object_store_for_path( parameter is provided. Args: - path: A URL-style path (e.g. ``"s3://bucket/key.parquet"``). + path: A URL-style path (e.g. ``"s3://bucket/key.parquet"`` or + ``"file:///tmp/data.parquet"``). store: An object store instance to register. Raises: - ValueError: If the path does not contain a recognized URL scheme. + ValueError: If the path does not contain a URL scheme, or if + a non-file scheme is missing a host/bucket component. """ parsed = urlparse(str(path)) - if not parsed.scheme or not parsed.netloc: + if not parsed.scheme: 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) + # file:// URLs typically have an empty netloc (e.g. file:///tmp/a.parquet) + # For other schemes (s3, gs, az, https) the netloc (bucket/host) is required. + if parsed.scheme != "file" and not parsed.netloc: + msg = ( + f"Cannot determine object store URL from path {path!r}. " + "The path must include a host or bucket " + "(e.g. 's3://bucket/key')." + ) + raise ValueError(msg) scheme = f"{parsed.scheme}://" - host = parsed.netloc + host = parsed.netloc or None self.register_object_store(scheme, store, host=host) def register_listing_table( @@ -1096,7 +1107,7 @@ def register_parquet( >>> import datafusion >>> ctx = datafusion.SessionContext() - >>> ctx.register_parquet("my_table", "data.parquet") + >>> ctx.register_parquet("my_table", "data.parquet") # doctest: +SKIP Register from S3 with inline credentials (thread-safe): diff --git a/python/tests/test_object_store_param.py b/python/tests/test_object_store_param.py index b57d9f287..42ac003ac 100644 --- a/python/tests/test_object_store_param.py +++ b/python/tests/test_object_store_param.py @@ -84,11 +84,26 @@ def test_raises_on_relative_path(self, ctx): def test_raises_on_windows_path(self, ctx): mock_store = MagicMock() - with pytest.raises(ValueError, match="Cannot determine object store URL"): + with pytest.raises(ValueError, match="must include a host or bucket"): ctx._register_object_store_for_path( "C:\\Users\\data\\file.parquet", mock_store ) + def test_raises_on_scheme_without_host_for_non_file(self, ctx): + """Non-file schemes (s3, gs, etc.) require a host/bucket.""" + mock_store = MagicMock() + with pytest.raises(ValueError, match="must include a host or bucket"): + ctx._register_object_store_for_path("s3:///key.parquet", mock_store) + + def test_parses_file_url(self, ctx): + """file:// URLs with empty netloc should be accepted.""" + mock_store = MagicMock() + with patch.object(ctx, "register_object_store") as mock_register: + ctx._register_object_store_for_path( + "file:///tmp/path/to/file.parquet", mock_store + ) + mock_register.assert_called_once_with("file://", mock_store, host=None) + def test_accepts_pathlib_path_raises(self, ctx): """pathlib.Path cannot represent URLs, so this should raise.""" mock_store = MagicMock() @@ -258,7 +273,7 @@ def test_register_parquet_with_local_object_store(self, ctx, tmp_path): # Use file:// URL with LocalFileSystem object store store = LocalFileSystem() - file_url = f"file://{tmp_path}/test.parquet" + file_url = parquet_path.as_uri() ctx.register_parquet("test_tbl", file_url, object_store=store) result = ctx.sql("SELECT * FROM test_tbl").collect() @@ -277,7 +292,7 @@ def test_read_parquet_with_local_object_store(self, ctx, tmp_path): pq.write_table(table, str(parquet_path)) store = LocalFileSystem() - file_url = f"file://{tmp_path}/read_test.parquet" + file_url = parquet_path.as_uri() df = ctx.read_parquet(file_url, object_store=store) result = df.collect() From f8733a50f6b87aed9b80d55ddd38a97989709a78 Mon Sep 17 00:00:00 2001 From: Jared Yu Date: Thu, 30 Jul 2026 11:28:42 -0700 Subject: [PATCH 3/3] Refactor tests to use pytest.mark.parametrize per reviewer feedback Replace class-based test structure with flat parametrized functions to match the repository's testing conventions. --- python/tests/test_object_store_param.py | 375 +++++++----------------- 1 file changed, 107 insertions(+), 268 deletions(-) diff --git a/python/tests/test_object_store_param.py b/python/tests/test_object_store_param.py index 42ac003ac..13d51d792 100644 --- a/python/tests/test_object_store_param.py +++ b/python/tests/test_object_store_param.py @@ -22,8 +22,10 @@ from unittest.mock import MagicMock, patch import pyarrow as pa +import pyarrow.parquet as pq import pytest from datafusion import SessionContext +from datafusion.object_store import LocalFileSystem @pytest.fixture @@ -31,271 +33,108 @@ 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="must include a host or bucket"): - ctx._register_object_store_for_path( - "C:\\Users\\data\\file.parquet", mock_store - ) - - def test_raises_on_scheme_without_host_for_non_file(self, ctx): - """Non-file schemes (s3, gs, etc.) require a host/bucket.""" - mock_store = MagicMock() - with pytest.raises(ValueError, match="must include a host or bucket"): - ctx._register_object_store_for_path("s3:///key.parquet", mock_store) - - def test_parses_file_url(self, ctx): - """file:// URLs with empty netloc should be accepted.""" - mock_store = MagicMock() - with patch.object(ctx, "register_object_store") as mock_register: - ctx._register_object_store_for_path( - "file:///tmp/path/to/file.parquet", mock_store - ) - mock_register.assert_called_once_with("file://", mock_store, host=None) - - 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 = parquet_path.as_uri() - - ctx.register_parquet("test_tbl", file_url, object_store=store) - result = ctx.sql("SELECT * FROM test_tbl").collect() - - assert len(result) == 1 - assert result[0].num_rows == 3 - assert result[0].column("x").to_pylist() == [1, 2, 3] - - def test_read_parquet_with_local_object_store(self, ctx, tmp_path): - """Verify read_parquet works with object_store parameter.""" - import pyarrow.parquet as pq - from datafusion.object_store import LocalFileSystem - - table = pa.table({"val": [10, 20, 30]}) - parquet_path = tmp_path / "read_test.parquet" - pq.write_table(table, str(parquet_path)) - - store = LocalFileSystem() - file_url = parquet_path.as_uri() - - df = ctx.read_parquet(file_url, object_store=store) - result = df.collect() - - assert len(result) == 1 - assert result[0].column("val").to_pylist() == [10, 20, 30] +@pytest.mark.parametrize( + ("path", "scheme", "host"), + [ + ("s3://my-bucket/path/file.parquet", "s3://", "my-bucket"), + ("gs://my-gcs-bucket/data.parquet", "gs://", "my-gcs-bucket"), + ("az://my-container/data.parquet", "az://", "my-container"), + ("https://example.com/data.parquet", "https://", "example.com"), + ("file:///tmp/data.parquet", "file://", None), + ], +) +def test_register_object_store_for_url(ctx, path, scheme, host): + store = MagicMock() + + with patch.object(ctx, "register_object_store") as register: + ctx._register_object_store_for_path(path, store) + + register.assert_called_once_with(scheme, store, host=host) + + +@pytest.mark.parametrize( + ("path", "error"), + [ + ("/local/path/file.parquet", "Cannot determine object store URL"), + ("relative/path.parquet", "Cannot determine object store URL"), + (Path("/local/file.parquet"), "Cannot determine object store URL"), + ("C:\\Users\\data\\file.parquet", "must include a host or bucket"), + ("s3:///key.parquet", "must include a host or bucket"), + ], +) +def test_register_object_store_rejects_invalid_url(path, error, ctx): + with pytest.raises(ValueError, match=error): + ctx._register_object_store_for_path(path, MagicMock()) + + +@pytest.mark.parametrize( + ("method_name", "args", "path"), + [ + ("register_parquet", ("table",), "s3://bucket/data.parquet"), + ("read_parquet", (), "s3://bucket/data.parquet"), + ("register_csv", ("table",), "s3://bucket/data.csv"), + ("read_csv", (), "s3://bucket/data.csv"), + ("register_json", ("table",), "s3://bucket/data.json"), + ("read_json", (), "s3://bucket/data.json"), + ("register_avro", ("table",), "s3://bucket/data.avro"), + ("read_avro", (), "s3://bucket/data.avro"), + ("register_arrow", ("table",), "s3://bucket/data.arrow"), + ("read_arrow", (), "s3://bucket/data.arrow"), + ], +) +def test_file_methods_register_object_store(ctx, method_name, args, path): + store = MagicMock() + + # The remote file does not exist. Registration happens before DataFusion + # tries to inspect it, which is the behavior under test. + with ( + patch.object(ctx, "register_object_store") as register, + contextlib.suppress(Exception), + ): + getattr(ctx, method_name)(*args, path, object_store=store) + + register.assert_called_once_with("s3://", store, host="bucket") + + +def test_register_csv_uses_first_path_for_object_store(ctx): + store = MagicMock() + paths = ["s3://bucket/a.csv", "s3://bucket/b.csv"] + + with ( + patch.object(ctx, "register_object_store") as register, + contextlib.suppress(Exception), + ): + ctx.register_csv("table", paths, object_store=store) + + register.assert_called_once_with("s3://", store, host="bucket") + + +def test_object_store_none_does_not_register(ctx): + with ( + patch.object(ctx, "register_object_store") as register, + contextlib.suppress(Exception), + ): + ctx.register_parquet("table", "missing.parquet") + + register.assert_not_called() + + +def test_file_method_rejects_local_path_with_object_store(ctx): + with pytest.raises(ValueError, match="Cannot determine object store URL"): + ctx.register_parquet("table", "/local/file.parquet", object_store=MagicMock()) + + +@pytest.mark.parametrize("method_name", ["register_parquet", "read_parquet"]) +def test_parquet_methods_with_local_object_store(ctx, tmp_path, method_name): + table = pa.table({"value": [10, 20, 30]}) + parquet_path = tmp_path / "data.parquet" + pq.write_table(table, parquet_path) + + path = parquet_path.as_uri() + if method_name == "register_parquet": + ctx.register_parquet("test_table", path, object_store=LocalFileSystem()) + dataframe = ctx.sql("SELECT * FROM test_table") + else: + dataframe = ctx.read_parquet(path, object_store=LocalFileSystem()) + + assert dataframe.collect()[0].column("value").to_pylist() == [10, 20, 30]