Add object_store parameter to register/read file methods for thread-safe credentials (#1625)

* 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

* 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

* 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.

---------

Co-authored-by: Tim Saucer <timsaucer@gmail.com>

Author: qzyu999

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,45 @@ 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"`` or
+                ``"file:///tmp/data.parquet"``).
+            store: An object store instance to register.
+
+        Raises:
+            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:
+            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 or None
+        self.register_object_store(scheme, store, host=host)
+
     def register_listing_table(
         self,
         name: str,
@@ -1028,6 +1069,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 +1091,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")  # doctest: +SKIP
+
+            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 +1151,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 +1173,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 +1227,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 +1244,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 +1270,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 +1283,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 +1301,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 +1314,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 +1371,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 +1792,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 +1805,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 +1839,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 +1859,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 +1918,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 +1938,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 +1997,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 +2006,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 +2028,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 +2037,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 +2091,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)