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]