atlan-application-sdk 0.1.1rc41__py3-none-any.whl → 0.1.1rc42__py3-none-any.whl

application_sdk/activities/common/utils.py

@@ -5,10 +5,11 @@ including workflow ID retrieval, automatic heartbeating, and periodic heartbeat
 """
 
 import asyncio
+import glob
 import os
 from datetime import timedelta
 from functools import wraps
-from typing import Any, Awaitable, Callable, Optional, TypeVar, cast
+from typing import Any, Awaitable, Callable, List, Optional, TypeVar, cast
 
 from temporalio import activity
 
@@ -79,17 +80,47 @@ def build_output_path() -> str:
 
 def get_object_store_prefix(path: str) -> str:
     """Get the object store prefix for the path.
+
+    This function handles two types of paths:
+    1. Paths under TEMPORARY_PATH - converts them to relative object store paths
+    2. User-provided paths - returns them as-is (already relative object store paths)
+
     Args:
-        path: The path to the output directory.
+        path: The path to convert to object store prefix.
 
     Returns:
         The object store prefix for the path.
 
-    Example:
+    Examples:
+        >>> # Temporary path case
         >>> get_object_store_prefix("./local/tmp/artifacts/apps/appName/workflows/wf-123/run-456")
         "artifacts/apps/appName/workflows/wf-123/run-456"
+
+        >>> # User-provided path case
+        >>> get_object_store_prefix("datasets/sales/2024/")
+        "datasets/sales/2024"
     """
-    return os.path.relpath(path, TEMPORARY_PATH)
+    # Normalize paths for comparison
+    abs_path = os.path.abspath(path)
+    abs_temp_path = os.path.abspath(TEMPORARY_PATH)
+
+    # Check if path is under TEMPORARY_PATH
+    try:
+        # Use os.path.commonpath to properly check if path is under temp directory
+        # This prevents false positives like '/tmp/local123' matching '/tmp/local'
+        common_path = os.path.commonpath([abs_path, abs_temp_path])
+        if common_path == abs_temp_path:
+            # Path is under temp directory, convert to relative object store path
+            relative_path = os.path.relpath(abs_path, abs_temp_path)
+            # Normalize path separators to forward slashes for object store
+            return relative_path.replace(os.path.sep, "/")
+        else:
+            # Path is already a relative object store path, return as-is
+            return path.strip("/")
+    except ValueError:
+        # os.path.commonpath or os.path.relpath can raise ValueError on Windows with different drives
+        # In this case, treat as user-provided path, return as-is
+        return path.strip("/")
 
 
 def auto_heartbeater(fn: F) -> F:
@@ -199,3 +230,46 @@ async def send_periodic_heartbeat(delay: float, *details: Any) -> None:
     while True:
         await asyncio.sleep(delay)
         activity.heartbeat(*details)
+
+
+def find_local_files_by_extension(
+    path: str,
+    extension: str,
+    file_names: Optional[List[str]] = None,
+) -> List[str]:
+    """Find local files at the specified local path, optionally filtering by file names.
+
+    Args:
+        path (str): Local path to search in (file or directory)
+        extension (str): File extension to filter by (e.g., '.parquet', '.json')
+        file_names (Optional[List[str]]): List of file names (basenames) to filter by, paths are not supported
+
+    Returns:
+        List[str]: List of matching file paths
+
+    Example:
+        >>> find_local_files_by_extension("/data", ".parquet", ["file1.parquet", "file2.parquet"])
+        ['file1.parquet', 'file2.parquet']
+
+        >>> find_local_files_by_extension("/data/single.json", ".json")
+        ['single.json']
+    """
+    if os.path.isfile(path) and path.endswith(extension):
+        # Single file - return it directly
+        return [path]
+
+    elif os.path.isdir(path):
+        # Directory - find all files in directory
+        all_files = glob.glob(
+            os.path.join(path, "**", f"*{extension}"),
+            recursive=True,
+        )
+
+        # Filter by file names if specified
+        if file_names:
+            file_names_set = set(file_names)  # Convert to set for O(1) lookup
+            return [f for f in all_files if os.path.basename(f) in file_names_set]
+        else:
+            return all_files
+
+    return []

application_sdk/activities/metadata_extraction/sql.py

@@ -860,9 +860,7 @@ class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
 
         raw_input = ParquetInput(
             path=os.path.join(output_path, "raw"),
-            input_prefix=output_prefix,
             file_names=workflow_args.get("file_names"),
-            chunk_size=None,
         )
         raw_input = raw_input.get_batched_daft_dataframe()
         transformed_output = JsonOutput(

application_sdk/inputs/__init__.py

@@ -1,7 +1,15 @@
+import os
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, AsyncIterator, Iterator, Union
-
+from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Union
+
+from application_sdk.activities.common.utils import (
+    find_local_files_by_extension,
+    get_object_store_prefix,
+)
+from application_sdk.common.error_codes import IOError
+from application_sdk.constants import TEMPORARY_PATH
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 
@@ -15,6 +23,94 @@ class Input(ABC):
     Abstract base class for input data sources.
     """
 
+    async def download_files(self) -> List[str]:
+        """Download files from object store if not available locally.
+
+        Flow:
+        1. Check if files exist locally at self.path
+        2. If not, try to download from object store
+        3. Filter by self.file_names if provided
+        4. Return list of file paths for logging purposes
+
+        Returns:
+            List[str]: List of file paths
+
+        Raises:
+            AttributeError: When the input class doesn't support file operations or _extension
+            IOError: When no files found locally or in object store
+        """
+        # Step 1: Check if files exist locally
+        local_files = find_local_files_by_extension(
+            self.path, self._EXTENSION, self.file_names
+        )
+        if local_files:
+            logger.info(
+                f"Found {len(local_files)} {self._EXTENSION} files locally at: {self.path}"
+            )
+            return local_files
+
+        # Step 2: Try to download from object store
+        logger.info(
+            f"No local {self._EXTENSION} files found at {self.path}, checking object store..."
+        )
+
+        try:
+            # Determine what to download based on path type and filters
+            downloaded_paths = []
+
+            if self.path.endswith(self._EXTENSION):
+                # Single file case (file_names validation already ensures this is valid)
+                source_path = get_object_store_prefix(self.path)
+                destination_path = os.path.join(TEMPORARY_PATH, source_path)
+                await ObjectStore.download_file(
+                    source=source_path,
+                    destination=destination_path,
+                )
+                downloaded_paths.append(destination_path)
+
+            elif self.file_names:
+                # Directory with specific files - download each file individually
+                for file_name in self.file_names:
+                    file_path = os.path.join(self.path, file_name)
+                    source_path = get_object_store_prefix(file_path)
+                    destination_path = os.path.join(TEMPORARY_PATH, source_path)
+                    await ObjectStore.download_file(
+                        source=source_path,
+                        destination=destination_path,
+                    )
+                    downloaded_paths.append(destination_path)
+            else:
+                # Download entire directory
+                source_path = get_object_store_prefix(self.path)
+                destination_path = os.path.join(TEMPORARY_PATH, source_path)
+                await ObjectStore.download_prefix(
+                    source=source_path,
+                    destination=destination_path,
+                )
+                # Find the actual files in the downloaded directory
+                found_files = find_local_files_by_extension(
+                    destination_path, self._EXTENSION, getattr(self, "file_names", None)
+                )
+                downloaded_paths.extend(found_files)
+
+            # Check results
+            if downloaded_paths:
+                logger.info(
+                    f"Successfully downloaded {len(downloaded_paths)} {self._EXTENSION} files from object store"
+                )
+                return downloaded_paths
+            else:
+                raise IOError(
+                    f"{IOError.OBJECT_STORE_READ_ERROR}: Downloaded from object store but no {self._EXTENSION} files found"
+                )
+
+        except Exception as e:
+            logger.error(f"Failed to download from object store: {str(e)}")
+            raise IOError(
+                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: No {self._EXTENSION} files found locally at '{self.path}' and failed to download from object store. "
+                f"Error: {str(e)}"
+            )
+
     @abstractmethod
     async def get_batched_dataframe(
         self,

application_sdk/inputs/json.py

@@ -1,11 +1,7 @@
-import os
 from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Union
 
-from application_sdk.activities.common.utils import get_object_store_prefix
-from application_sdk.common.error_codes import IOError
 from application_sdk.inputs import Input
 from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.services.objectstore import ObjectStore
 
 if TYPE_CHECKING:
     import daft
@@ -15,56 +11,43 @@ logger = get_logger(__name__)
 
 
 class JsonInput(Input):
-    path: str
-    chunk_size: Optional[int]
-    file_names: Optional[List[str]]
-    download_file_prefix: Optional[str]
+    """
+    JSON Input class to read data from JSON files using daft and pandas.
+    Supports reading both single files and directories containing multiple JSON files.
+    """
+
+    _EXTENSION = ".json"
 
     def __init__(
         self,
         path: str,
         file_names: Optional[List[str]] = None,
-        download_file_prefix: Optional[str] = None,
-        chunk_size: Optional[int] = None,
+        chunk_size: int = 100000,
     ):
         """Initialize the JsonInput class.
 
         Args:
-            path (str): The path to the input directory.
-            file_names (Optional[List[str]]): The list of files to read.
-            download_file_prefix (Optional[str]): The prefix path in object store.
-            chunk_size (Optional[int]): The chunk size to read the data. If None, uses config value.
+            path (str): Path to JSON file or directory containing JSON files.
+                It accepts both types of paths:
+                local path or object store path
+                Wildcards are not supported.
+            file_names (Optional[List[str]]): List of specific file names to read. Defaults to None.
+            chunk_size (int): Number of rows per batch. Defaults to 100000.
+
+        Raises:
+            ValueError: When path is not provided or when single file path is combined with file_names
         """
+
+        # Validate that single file path and file_names are not both specified
+        if path.endswith(self._EXTENSION) and file_names:
+            raise ValueError(
+                f"Cannot specify both a single file path ('{path}') and file_names filter. "
+                f"Either provide a directory path with file_names, or specify the exact file path without file_names."
+            )
+
         self.path = path
-        # If chunk_size is provided, use it; otherwise default to 100,000 rows per batch
-        self.chunk_size = chunk_size if chunk_size is not None else 100000
+        self.chunk_size = chunk_size
         self.file_names = file_names
-        self.download_file_prefix = download_file_prefix
-
-    async def download_files(self):
-        """Download the files from the object store to the local path"""
-        if not self.file_names:
-            logger.debug("No files to download")
-            return
-
-        for file_name in self.file_names or []:
-            try:
-                if self.download_file_prefix is not None and not os.path.exists(
-                    os.path.join(self.path, file_name)
-                ):
-                    destination_file_path = os.path.join(self.path, file_name)
-                    await ObjectStore.download_file(
-                        source=get_object_store_prefix(destination_file_path),
-                        destination=destination_file_path,
-                    )
-            except IOError as e:
-                logger.error(
-                    f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error downloading file {file_name}: {str(e)}",
-                    error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
-                )
-                raise IOError(
-                    f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error downloading file {file_name}: {str(e)}"
-                )
 
     async def get_batched_dataframe(
         self,
@@ -76,22 +59,20 @@ class JsonInput(Input):
         try:
             import pandas as pd
 
-            await self.download_files()
+            # Ensure files are available (local or downloaded)
+            json_files = await self.download_files()
+            logger.info(f"Reading {len(json_files)} JSON files in batches")
 
-            for file_name in self.file_names or []:
-                file_path = os.path.join(self.path, file_name)
+            for json_file in json_files:
                 json_reader_obj = pd.read_json(
-                    file_path,
+                    json_file,
                     chunksize=self.chunk_size,
                     lines=True,
                 )
                 for chunk in json_reader_obj:
                     yield chunk
-        except IOError as e:
-            logger.error(
-                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error reading batched data from JSON: {str(e)}",
-                error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
-            )
+        except Exception as e:
+            logger.error(f"Error reading batched data from JSON: {str(e)}")
             raise
 
     async def get_dataframe(self) -> "pd.DataFrame":
@@ -102,21 +83,17 @@ class JsonInput(Input):
         try:
             import pandas as pd
 
-            dataframes = []
-            await self.download_files()
-            for file_name in self.file_names or []:
-                dataframes.append(
-                    pd.read_json(
-                        os.path.join(self.path, file_name),
-                        lines=True,
-                    )
-                )
-            return pd.concat(dataframes, ignore_index=True)
-        except IOError as e:
-            logger.error(
-                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error reading data from JSON: {str(e)}",
-                error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
+            # Ensure files are available (local or downloaded)
+            json_files = await self.download_files()
+            logger.info(f"Reading {len(json_files)} JSON files as pandas dataframe")
+
+            return pd.concat(
+                (pd.read_json(json_file, lines=True) for json_file in json_files),
+                ignore_index=True,
             )
+
+        except Exception as e:
+            logger.error(f"Error reading data from JSON: {str(e)}")
             raise
 
     async def get_batched_daft_dataframe(
@@ -129,18 +106,15 @@ class JsonInput(Input):
         try:
             import daft
 
-            await self.download_files()
-            for file_name in self.file_names or []:
-                json_reader_obj = daft.read_json(
-                    path=os.path.join(self.path, file_name),
-                    _chunk_size=self.chunk_size,
-                )
-                yield json_reader_obj
-        except IOError as e:
-            logger.error(
-                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error reading batched data from JSON: {str(e)}",
-                error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
-            )
+            # Ensure files are available (local or downloaded)
+            json_files = await self.download_files()
+            logger.info(f"Reading {len(json_files)} JSON files as daft batches")
+
+            # Yield each discovered file as separate batch with chunking
+            for json_file in json_files:
+                yield daft.read_json(json_file, _chunk_size=self.chunk_size)
+        except Exception as e:
+            logger.error(f"Error reading batched data from JSON using daft: {str(e)}")
             raise
 
     async def get_daft_dataframe(self) -> "daft.DataFrame":  # noqa: F821
@@ -151,14 +125,12 @@ class JsonInput(Input):
         try:
             import daft
 
-            await self.download_files()
-            if not self.file_names or len(self.file_names) == 0:
-                raise ValueError("No files to read")
-            directory = os.path.join(self.path, self.file_names[0].split("/")[0])
-            return daft.read_json(path=f"{directory}/*.json")
-        except IOError as e:
-            logger.error(
-                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error reading data from JSON using daft: {str(e)}",
-                error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
-            )
+            # Ensure files are available (local or downloaded)
+            json_files = await self.download_files()
+            logger.info(f"Reading {len(json_files)} JSON files with daft")
+
+            # Use the discovered/downloaded files directly
+            return daft.read_json(json_files)
+        except Exception as e:
+            logger.error(f"Error reading data from JSON using daft: {str(e)}")
             raise

application_sdk/inputs/parquet.py

@@ -1,11 +1,7 @@
-import glob
-import os
 from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Union
 
-from application_sdk.activities.common.utils import get_object_store_prefix
 from application_sdk.inputs import Input
 from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 
@@ -20,107 +16,139 @@ class ParquetInput(Input):
     Supports reading both single files and directories containing multiple parquet files.
     """
 
+    _EXTENSION = ".parquet"
+
     def __init__(
         self,
-        path: Optional[str] = None,
-        chunk_size: Optional[int] = 100000,
-        input_prefix: Optional[str] = None,
+        path: str,
+        chunk_size: int = 100000,
         file_names: Optional[List[str]] = None,
     ):
         """Initialize the Parquet input class.
 
         Args:
             path (str): Path to parquet file or directory containing parquet files.
-            chunk_size (Optional[int], optional): Number of rows per batch.
-                Defaults to 100000.
-            input_prefix (Optional[str], optional): Prefix for files when reading from object store.
-                If provided, files will be read from object store. Defaults to None.
-            file_names (Optional[List[str]], optional): List of file names to read.
-                Defaults to None.
+                It accepts both types of paths:
+                local path or object store path
+                Wildcards are not supported.
+            chunk_size (int): Number of rows per batch. Defaults to 100000.
+            file_names (Optional[List[str]]): List of file names to read. Defaults to None.
+
+        Raises:
+            ValueError: When path is not provided or when single file path is combined with file_names
         """
+
+        # Validate that single file path and file_names are not both specified
+        if path.endswith(self._EXTENSION) and file_names:
+            raise ValueError(
+                f"Cannot specify both a single file path ('{path}') and file_names filter. "
+                f"Either provide a directory path with file_names, or specify the exact file path without file_names."
+            )
+
         self.path = path
         self.chunk_size = chunk_size
-        self.input_prefix = input_prefix
         self.file_names = file_names
 
-    async def download_files(self, local_path: str) -> Optional[str]:
-        """Read a file from the object store.
-
-        Args:
-            local_path (str): Path to the local data in the temp directory.
+    async def get_dataframe(self) -> "pd.DataFrame":
+        """Read data from parquet file(s) and return as pandas DataFrame.
 
         Returns:
-            Optional[str]: Path to the downloaded local file.
-        """
-        # if the path is a directory, then check if the directory has any parquet files
-        parquet_files = []
-        if os.path.isdir(local_path):
-            parquet_files = glob.glob(os.path.join(local_path, "*.parquet"))
-        else:
-            parquet_files = glob.glob(local_path)
-        if not parquet_files:
-            if self.input_prefix:
-                logger.info(
-                    f"Reading file from object store: {local_path} from {self.input_prefix}"
-                )
-                if os.path.isdir(local_path):
-                    await ObjectStore.download_prefix(
-                        source=get_object_store_prefix(local_path),
-                        destination=local_path,
-                    )
-                else:
-                    await ObjectStore.download_file(
-                        source=get_object_store_prefix(local_path),
-                        destination=local_path,
-                    )
-            else:
-                raise ValueError(
-                    f"No parquet files found in {local_path} and no input prefix provided"
-                )
+            pd.DataFrame: Combined dataframe from specified parquet files
 
-    async def get_dataframe(self) -> "pd.DataFrame":
-        """
-        Method to read the data from the parquet file(s)
-        and return as a single combined pandas dataframe.
+        Raises:
+            ValueError: When no valid path can be determined or no matching files found
+            Exception: When reading parquet files fails
 
-        Returns:
-            "pd.DataFrame": Combined dataframe from all parquet files.
+        Example transformation:
+        Input files:
+        +------------------+
+        | file1.parquet    |
+        | file2.parquet    |
+        | file3.parquet    |
+        +------------------+
+
+        With file_names=["file1.parquet", "file3.parquet"]:
+        +-------+-------+-------+
+        | col1  | col2  | col3  |
+        +-------+-------+-------+
+        | val1  | val2  | val3  |  # from file1.parquet
+        | val7  | val8  | val9  |  # from file3.parquet
+        +-------+-------+-------+
+
+        Transformations:
+        - Only specified files are read and combined
+        - Column schemas must be compatible across files
+        - Only reads files in the specified directory
         """
         try:
             import pandas as pd
 
-            path = self.path
-            if self.input_prefix and self.path:
-                path = await self.download_files(self.path)
-            # Use pandas native read_parquet which can handle both single files and directories
-            return pd.read_parquet(path)
+            # Ensure files are available (local or downloaded)
+            parquet_files = await self.download_files()
+            logger.info(f"Reading {len(parquet_files)} parquet files")
+
+            return pd.concat(
+                (pd.read_parquet(parquet_file) for parquet_file in parquet_files),
+                ignore_index=True,
+            )
         except Exception as e:
             logger.error(f"Error reading data from parquet file(s): {str(e)}")
-            # Re-raise to match IcebergInput behavior
             raise
 
     async def get_batched_dataframe(
         self,
     ) -> Union[AsyncIterator["pd.DataFrame"], Iterator["pd.DataFrame"]]:
-        """
-        Method to read the data from the parquet file(s) in batches
-        and return as an async iterator of pandas dataframes.
+        """Read data from parquet file(s) in batches as pandas DataFrames.
 
         Returns:
-            AsyncIterator["pd.DataFrame"]: Async iterator of pandas dataframes.
+            AsyncIterator[pd.DataFrame]: Async iterator of pandas dataframes
+
+        Raises:
+            ValueError: When no parquet files found locally or in object store
+            Exception: When reading parquet files fails
+
+        Example transformation:
+        Input files:
+        +------------------+
+        | file1.parquet    |
+        | file2.parquet    |
+        | file3.parquet    |
+        +------------------+
+
+        With file_names=["file1.parquet", "file2.parquet"] and chunk_size=2:
+        Batch 1:
+        +-------+-------+
+        | col1  | col2  |
+        +-------+-------+
+        | val1  | val2  |  # from file1.parquet
+        | val3  | val4  |  # from file1.parquet
+        +-------+-------+
+
+        Batch 2:
+        +-------+-------+
+        | col1  | col2  |
+        +-------+-------+
+        | val5  | val6  |  # from file2.parquet
+        | val7  | val8  |  # from file2.parquet
+        +-------+-------+
+
+        Transformations:
+        - Only specified files are combined then split into chunks
+        - Each batch is a separate DataFrame
+        - Only reads files in the specified directory
         """
         try:
             import pandas as pd
 
-            path = self.path
-            if self.input_prefix and self.path:
-                path = await self.download_files(self.path)
-            df = pd.read_parquet(path)
-            if self.chunk_size:
+            # Ensure files are available (local or downloaded)
+            parquet_files = await self.download_files()
+            logger.info(f"Reading {len(parquet_files)} parquet files in batches")
+
+            # Process each file individually to maintain memory efficiency
+            for parquet_file in parquet_files:
+                df = pd.read_parquet(parquet_file)
                 for i in range(0, len(df), self.chunk_size):
                     yield df.iloc[i : i + self.chunk_size]
-            else:
-                yield df
         except Exception as e:
             logger.error(
                 f"Error reading data from parquet file(s) in batches: {str(e)}"
@@ -128,51 +156,102 @@ class ParquetInput(Input):
             raise
 
     async def get_daft_dataframe(self) -> "daft.DataFrame":  # noqa: F821
-        """
-        Method to read the data from the parquet file(s)
-        and return as a single combined daft dataframe.
+        """Read data from parquet file(s) and return as daft DataFrame.
 
         Returns:
-            daft.DataFrame: Combined daft dataframe from all parquet files.
+            daft.DataFrame: Combined daft dataframe from specified parquet files
+
+        Raises:
+            ValueError: When no parquet files found locally or in object store
+            Exception: When reading parquet files fails
+
+        Example transformation:
+        Input files:
+        +------------------+
+        | file1.parquet    |
+        | file2.parquet    |
+        | file3.parquet    |
+        +------------------+
+
+        With file_names=["file1.parquet", "file3.parquet"]:
+        +-------+-------+-------+
+        | col1  | col2  | col3  |
+        +-------+-------+-------+
+        | val1  | val2  | val3  |  # from file1.parquet
+        | val7  | val8  | val9  |  # from file3.parquet
+        +-------+-------+-------+
+
+        Transformations:
+        - Only specified parquet files combined into single daft DataFrame
+        - Lazy evaluation for better performance
+        - Column schemas must be compatible across files
         """
         try:
             import daft  # type: ignore
 
-            if self.file_names:
-                path = f"{self.path}/{self.file_names[0].split('/')[0]}"
-            else:
-                path = self.path
-            if self.input_prefix and path:
-                await self.download_files(path)
-            return daft.read_parquet(f"{path}/*.parquet")
+            # Ensure files are available (local or downloaded)
+            parquet_files = await self.download_files()
+            logger.info(f"Reading {len(parquet_files)} parquet files with daft")
+
+            # Use the discovered/downloaded files directly
+            return daft.read_parquet(parquet_files)
         except Exception as e:
             logger.error(
                 f"Error reading data from parquet file(s) using daft: {str(e)}"
             )
-            # Re-raise to match IcebergInput behavior
             raise
 
     async def get_batched_daft_dataframe(self) -> AsyncIterator["daft.DataFrame"]:  # type: ignore
-        """
-        Get batched daft dataframe from parquet file(s)
+        """Get batched daft dataframe from parquet file(s).
 
         Returns:
             AsyncIterator[daft.DataFrame]: An async iterator of daft DataFrames, each containing
-            a batch of data from the parquet file(s).
+            a batch of data from individual parquet files
+
+        Raises:
+            ValueError: When no parquet files found locally or in object store
+            Exception: When reading parquet files fails
+
+        Example transformation:
+        Input files:
+        +------------------+
+        | file1.parquet    |
+        | file2.parquet    |
+        | file3.parquet    |
+        +------------------+
+
+        With file_names=["file1.parquet", "file3.parquet"]:
+        Batch 1 (file1.parquet):
+        +-------+-------+
+        | col1  | col2  |
+        +-------+-------+
+        | val1  | val2  |
+        | val3  | val4  |
+        +-------+-------+
+
+        Batch 2 (file3.parquet):
+        +-------+-------+
+        | col1  | col2  |
+        +-------+-------+
+        | val7  | val8  |
+        | val9  | val10 |
+        +-------+-------+
+
+        Transformations:
+        - Each specified file becomes a separate daft DataFrame batch
+        - Lazy evaluation for better performance
+        - Files processed individually for memory efficiency
         """
         try:
             import daft  # type: ignore
 
-            if self.file_names:
-                for file_name in self.file_names:
-                    path = f"{self.path}/{file_name}"
-                    if self.input_prefix and path:
-                        await self.download_files(path)
-                        yield daft.read_parquet(path)
-            else:
-                if self.path and self.input_prefix:
-                    await self.download_files(self.path)
-                    yield daft.read_parquet(f"{self.path}/*.parquet")
+            # Ensure files are available (local or downloaded)
+            parquet_files = await self.download_files()
+            logger.info(f"Reading {len(parquet_files)} parquet files as daft batches")
+
+            # Yield each discovered file as separate batch
+            for parquet_file in parquet_files:
+                yield daft.read_parquet(parquet_file)
 
         except Exception as error:
             logger.error(

application_sdk/test_utils/hypothesis/strategies/inputs/json_input.py

@@ -2,11 +2,17 @@ from hypothesis import strategies as st
 
 # Strategy for generating safe file path components
 safe_path_strategy = st.text(
-    alphabet=st.characters(),
-)
+    alphabet="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-",
+    min_size=1,
+    max_size=20,
+).map(lambda x: f"/data/{x}")
 
 # Strategy for generating file names
-file_name_strategy = st.builds(lambda name: f"{name}.json", name=safe_path_strategy)
+file_name_strategy = st.text(
+    alphabet="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-",
+    min_size=1,
+    max_size=10,
+).map(lambda x: f"{x}.json")
 
 # Strategy for generating lists of file names
 file_names_strategy = st.lists(file_name_strategy, unique=True)
@@ -18,7 +24,6 @@ download_prefix_strategy = safe_path_strategy
 json_input_config_strategy = st.fixed_dictionaries(
     {
         "path": safe_path_strategy,
-        "download_file_prefix": download_prefix_strategy,
-        "file_names": file_names_strategy,
+        "file_names": st.one_of(st.none(), file_names_strategy),
     }
 )

application_sdk/test_utils/hypothesis/strategies/inputs/parquet_input.py

@@ -2,11 +2,17 @@ from hypothesis import strategies as st
 
 # Strategy for generating safe file path components
 safe_path_strategy = st.text(
-    alphabet=st.characters(),
-)
+    alphabet="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-",
+    min_size=1,
+    max_size=20,
+).map(lambda x: f"/data/{x}")
 
 # Strategy for generating file names
-file_name_strategy = st.builds(lambda name: f"{name}.parquet", name=safe_path_strategy)
+file_name_strategy = st.text(
+    alphabet="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-",
+    min_size=1,
+    max_size=10,
+).map(lambda x: f"{x}.parquet")
 
 # Strategy for generating lists of file names
 file_names_strategy = st.lists(file_name_strategy, unique=True)
@@ -22,7 +28,6 @@ parquet_input_config_strategy = st.fixed_dictionaries(
     {
         "path": safe_path_strategy,
         "chunk_size": chunk_size_strategy,
-        "input_prefix": st.one_of(st.none(), input_prefix_strategy),
         "file_names": st.one_of(st.none(), file_names_strategy),
     }
 )

application_sdk/version.py

@@ -2,4 +2,4 @@
 Version information for the application_sdk package.
 """
 
-__version__ = "0.1.1rc41"
+__version__ = "0.1.1rc42"

{atlan_application_sdk-0.1.1rc41.dist-info → atlan_application_sdk-0.1.1rc42.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 0.1.1rc41
+Version: 0.1.1rc42
 Summary: Atlan Application SDK is a Python library for developing applications on the Atlan Platform
 Project-URL: Repository, https://github.com/atlanhq/application-sdk
 Project-URL: Documentation, https://github.com/atlanhq/application-sdk/README.md

{atlan_application_sdk-0.1.1rc41.dist-info → atlan_application_sdk-0.1.1rc42.dist-info}/RECORD

@@ -1,17 +1,17 @@
 application_sdk/__init__.py,sha256=2e2mvmLJ5dxmJGPELtb33xwP-j6JMdoIuqKycEn7hjg,151
 application_sdk/constants.py,sha256=1THiejjOEgm4kHFN-PrwrUkfRk7q1pjOLWLm-t2ph1Q,10674
-application_sdk/version.py,sha256=zV2HH6SOEd5nUl8rfci_44ydqpKFM-6ShVxYm8ggtCk,88
+application_sdk/version.py,sha256=Q8mQScZPLmTAD4YMM35CTgtJ-tDFJqdKLeeq94J4GkU,88
 application_sdk/worker.py,sha256=i5f0AeKI39IfsLO05QkwC6uMz0zDPSJqP7B2byri1VI,7489
 application_sdk/activities/__init__.py,sha256=QaXLOBYbb0zPOY5kfDQh56qbXQFaYNXOjJ5PCvatiZ4,9530
 application_sdk/activities/lock_management.py,sha256=L__GZ9BsArwU1ntYwAgCKsSjCqN6QBeOfT-OT4WyD4Y,3983
 application_sdk/activities/.cursor/BUGBOT.md,sha256=FNykX5aMkdOhzgpiGqstOnSp9JN63iR2XP3onU4AGh8,15843
 application_sdk/activities/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/activities/common/models.py,sha256=LIZfWvTtgtbAUvvn-rwrPQgD7fP2J0Gxdxr_ITgw-jM,1243
-application_sdk/activities/common/utils.py,sha256=F4Fq9Gl_gvUQj_fSdwzTU7obqUnemYL1dgb_yS34vTM,6967
+application_sdk/activities/common/utils.py,sha256=nSNGkY5eS5pPc8etdPWkXBFTSaConGAD8LDtNqOMHF4,9836
 application_sdk/activities/metadata_extraction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/activities/metadata_extraction/base.py,sha256=ENFojpxqKdN_eVSL4iet3cGfylPOfcl1jnflfo4zhs8,3920
 application_sdk/activities/metadata_extraction/rest.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-application_sdk/activities/metadata_extraction/sql.py,sha256=02u1KXT0EViqKRhYoJ-aYytohAu4IbmTe_JPyC3axTw,35898
+application_sdk/activities/metadata_extraction/sql.py,sha256=I6TfA_sRb9w6slBhXuqJtw_2_4YSyK-1MiCHb4NWf-E,35829
 application_sdk/activities/query_extraction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/activities/query_extraction/sql.py,sha256=mesGP_kiWzrJ8wboWFVt2jbDuGG_Fl3kQVvVMdH3KWA,21228
 application_sdk/application/__init__.py,sha256=PbSImXYaQQ2IIee2SvI8AjDiSo2QcCFrM1PX3x-_RQs,8035
@@ -57,10 +57,10 @@ application_sdk/events/models.py,sha256=7Esqp3WlbriT2EqT4kNiY_sHtRXRPLj27b8SbeC5
 application_sdk/handlers/__init__.py,sha256=3Wf7jCVFR2nYOyHZEc9jj8BQUnHCylFqoezp70J2Df0,1329
 application_sdk/handlers/base.py,sha256=ieWFbv8Gm7vfrrpS-mdMSm-mHGuQY02qiAVX2qPdj3w,2467
 application_sdk/handlers/sql.py,sha256=6A_9xCtkXyNY5gPhImbftzrdPIEWIeTTqjyIewVESHA,17815
-application_sdk/inputs/__init__.py,sha256=_d-cUhcDyoJTJR3PdQkC831go6VDw9AM6Bg7-qm3NHI,1900
+application_sdk/inputs/__init__.py,sha256=_O5lK2A5EYyqwid8txKNEds3pHkoHGKrSTTWnQ-UzRA,6022
 application_sdk/inputs/iceberg.py,sha256=xiv1kNtVx1k0h3ZJbJeXjZwdfBGSy9j9orYP_AyCYlI,2756
-application_sdk/inputs/json.py,sha256=Yv70Y9YuutN2trqK5-z2UNtBL0895ZbdEiBDt9cYM9s,6216
-application_sdk/inputs/parquet.py,sha256=GnyB0r4-7GNLBl3ooVFUzsxunZsrHStKK2h7XRc7AIY,6723
+application_sdk/inputs/json.py,sha256=ZOgB3tuZSsb2m_KxiAdnbUQgU5ythCs-Mq-n4pPfeHA,4905
+application_sdk/inputs/parquet.py,sha256=51Wyvbv8vS6T_3bKHgq6pCva8w3PKCDH5jDuENy0z8c,9060
 application_sdk/inputs/sql_query.py,sha256=1EREgea6kKNaMIyX2HLJgbJ07rtAgLasd9NyvDcdZok,10636
 application_sdk/inputs/.cursor/BUGBOT.md,sha256=hwKGDbopv3NU0bpC_ElpAPDFcS59GWS3TunObGC6eLQ,9731
 application_sdk/interceptors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -115,8 +115,8 @@ application_sdk/test_utils/hypothesis/strategies/handlers/sql/__init__.py,sha256
 application_sdk/test_utils/hypothesis/strategies/handlers/sql/sql_metadata.py,sha256=xnf62RyS4UzxW1kTDys_4mg3Avg7KRfqRdP6O81FHp4,1883
 application_sdk/test_utils/hypothesis/strategies/handlers/sql/sql_preflight.py,sha256=e9uo6Bx5w_ZAEu6bDTWbMbmzqB0MYl2dH-JlXg3bkV8,2648
 application_sdk/test_utils/hypothesis/strategies/inputs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-application_sdk/test_utils/hypothesis/strategies/inputs/json_input.py,sha256=9mPI585J4tfQTZiM_uoFXCq7qrLML7jwH5eDm_mzvm4,749
-application_sdk/test_utils/hypothesis/strategies/inputs/parquet_input.py,sha256=qOi5En011KiOWAusP6dCvzr1rtC_sc2IkWfj7hWvSh0,950
+application_sdk/test_utils/hypothesis/strategies/inputs/json_input.py,sha256=HfdCZnXIZFJiRuORpnmioXh8qHls9sWNSFDysy8il-o,913
+application_sdk/test_utils/hypothesis/strategies/inputs/parquet_input.py,sha256=agjRA9agpak_GmWiIt9bi_oLGvLM_eunxXfxcNHK3MQ,1081
 application_sdk/test_utils/hypothesis/strategies/outputs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/test_utils/hypothesis/strategies/outputs/json_output.py,sha256=p9wotUJwc-Wmm54_qVG5Ivp_mgl7YTeAcQfC6RXlxCc,1835
 application_sdk/test_utils/hypothesis/strategies/outputs/statestore.py,sha256=gmYBwePNoSI_pl2WTXOClgkruzRwkOX_1SmBaUTha0c,2903
@@ -152,8 +152,8 @@ application_sdk/workflows/metadata_extraction/__init__.py,sha256=jHUe_ZBQ66jx8bg
 application_sdk/workflows/metadata_extraction/sql.py,sha256=BhaZavEL8H3Jvf28FGcHtZwqdsUT_EHZ4VTqiaieWek,12278
 application_sdk/workflows/query_extraction/__init__.py,sha256=n066_CX5RpJz6DIxGMkKS3eGSRg03ilaCtsqfJWQb7Q,117
 application_sdk/workflows/query_extraction/sql.py,sha256=kT_JQkLCRZ44ZpaC4QvPL6DxnRIIVh8gYHLqRbMI-hA,4826
-atlan_application_sdk-0.1.1rc41.dist-info/METADATA,sha256=foqYmvYq45g_A7zojtZJ0F06AjTxMOeAEnBE9FkkMFE,5567
-atlan_application_sdk-0.1.1rc41.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-atlan_application_sdk-0.1.1rc41.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-atlan_application_sdk-0.1.1rc41.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
-atlan_application_sdk-0.1.1rc41.dist-info/RECORD,,
+atlan_application_sdk-0.1.1rc42.dist-info/METADATA,sha256=i1UmEi69_Uelfws3QsQ_AnUCMksuaeBb2O00kAXVGDc,5567
+atlan_application_sdk-0.1.1rc42.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+atlan_application_sdk-0.1.1rc42.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+atlan_application_sdk-0.1.1rc42.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
+atlan_application_sdk-0.1.1rc42.dist-info/RECORD,,