Fix import and add commeent

streaming/base/storage/download.py

@@ -30,6 +30,9 @@ def _create_s3_client(unsigned: bool = False) -> Any:
         Args:
             unsigned (bool, optional):  Set to True if it is a public bucket.
                 Defaults to ``False``.
+
+        Returns:
+            botocore.client.BaseClient: The S3 client.
         """
         if unsigned:
             # Client will be using unsigned mode in which public
@@ -335,8 +338,8 @@ def download_from_databricks_unity_catalog(remote: str, local: str, timeout: flo
         local (str): Local path (local filesystem).
         timeout (float): Timeout in seconds.
     """
-    from util import (create_databricks_unity_catalog_s3_client,
-                      get_databricks_unity_catalog_volume_info)
+    from streaming.base.storage.util import (create_databricks_unity_catalog_s3_client,
+                                             get_databricks_unity_catalog_volume_info)
 
     obj = urllib.parse.urlparse(remote)
     if obj.scheme != 'uc':

streaming/base/storage/util.py

@@ -14,6 +14,9 @@ def get_databricks_unity_catalog_volume_info(volume_name: str) -> Dict[str, Any]
 
     Args:
         volume_name (str): The name of the Databricks Unity Catalog volume.
+
+    Returns:
+        Dict[str, Any]: The volume info.
     """
     import requests
 
@@ -32,6 +35,9 @@ def create_databricks_unity_catalog_s3_client(volume_id: str, timeout: float) ->
     Args:
         volume_id (str): The ID of the Databricks Unity Catalog volume.
         timeout (float): The timeout for the S3 client.
+
+    Returns:
+        botocore.client.BaseClient: The S3 client.
     """
     import boto3
     from botocore.config import Config