atlan-application-sdk 1.1.1__py3-none-any.whl → 2.1.0__py3-none-any.whl

application_sdk/inputs/__init__.py

@@ -1,168 +0,0 @@
-import os
-from abc import ABC, abstractmethod
-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__)
-
-if TYPE_CHECKING:
-    import daft
-    import pandas as pd
-
-
-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,
-    ) -> Union[Iterator["pd.DataFrame"], AsyncIterator["pd.DataFrame"]]:
-        """
-        Get an iterator of batched pandas DataFrames.
-
-        Returns:
-            Iterator["pd.DataFrame"]: An iterator of batched pandas DataFrames.
-
-        Raises:
-            NotImplementedError: If the method is not implemented.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    async def get_dataframe(self) -> "pd.DataFrame":
-        """
-        Get a single pandas DataFrame.
-
-        Returns:
-            "pd.DataFrame": A pandas DataFrame.
-
-        Raises:
-            NotImplementedError: If the method is not implemented.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    async def get_batched_daft_dataframe(
-        self,
-    ) -> Union[Iterator["daft.DataFrame"], AsyncIterator["daft.DataFrame"]]:  # noqa: F821
-        """
-        Get an iterator of batched daft DataFrames.
-
-        Returns:
-            Iterator[daft.DataFrame]: An iterator of batched daft DataFrames.
-
-        Raises:
-            NotImplementedError: If the method is not implemented.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    async def get_daft_dataframe(self) -> "daft.DataFrame":  # noqa: F821
-        """
-        Get a single daft DataFrame.
-
-        Returns:
-            daft.DataFrame: A daft DataFrame.
-
-        Raises:
-            NotImplementedError: If the method is not implemented.
-        """
-        raise NotImplementedError

application_sdk/inputs/iceberg.py

@@ -1,75 +0,0 @@
-from typing import TYPE_CHECKING, AsyncIterator, Iterator, Optional, Union
-
-from pyiceberg.table import Table
-
-from application_sdk.inputs import Input
-from application_sdk.observability.logger_adaptor import get_logger
-
-logger = get_logger(__name__)
-
-if TYPE_CHECKING:
-    import daft
-    import pandas as pd
-
-
-class IcebergInput(Input):
-    """
-    Iceberg Input class to read data from Iceberg tables using daft and pandas
-    """
-
-    table: Table
-    chunk_size: Optional[int]
-
-    def __init__(self, table: Table, chunk_size: Optional[int] = 100000):
-        """Initialize the Iceberg input class.
-
-        Args:
-            table (Table): Iceberg table object.
-            chunk_size (Optional[int], optional): Number of rows per batch.
-                Defaults to 100000.
-        """
-        self.table = table
-        self.chunk_size = chunk_size
-
-    async def get_dataframe(self) -> "pd.DataFrame":
-        """
-        Method to read the data from the iceberg table
-        and return as a single combined pandas dataframe
-        """
-        try:
-            daft_dataframe = await self.get_daft_dataframe()
-            return daft_dataframe.to_pandas()
-        except Exception as e:
-            logger.error(f"Error reading data from Iceberg table: {str(e)}")
-            raise
-
-    async def get_batched_dataframe(
-        self,
-    ) -> Union[AsyncIterator["pd.DataFrame"], Iterator["pd.DataFrame"]]:
-        # We are not implementing this method as we have to partition the daft dataframe
-        # using dataframe.into_partitions() method. This method does all the partitions in memory
-        # and using that can cause out of memory issues.
-        # ref: https://www.getdaft.io/projects/docs/en/stable/user_guide/poweruser/partitioning.html
-        raise NotImplementedError
-
-    async def get_daft_dataframe(self) -> "daft.DataFrame":  # noqa: F821
-        """
-        Method to read the data from the iceberg table
-        and return as a single combined daft dataframe
-        """
-        try:
-            import daft
-
-            return daft.read_iceberg(self.table)
-        except Exception as e:
-            logger.error(f"Error reading data from Iceberg table using daft: {str(e)}")
-            raise
-
-    async def get_batched_daft_dataframe(
-        self,
-    ) -> Union[AsyncIterator["daft.DataFrame"], Iterator["daft.DataFrame"]]:  # noqa: F821
-        # We are not implementing this method as we have to partition the daft dataframe
-        # using dataframe.into_partitions() method. This method does all the partitions in memory
-        # and using that can cause out of memory issues.
-        # ref: https://www.getdaft.io/projects/docs/en/stable/user_guide/poweruser/partitioning.html
-        raise NotImplementedError

application_sdk/inputs/json.py

@@ -1,136 +0,0 @@
-from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Union
-
-from application_sdk.inputs import Input
-from application_sdk.observability.logger_adaptor import get_logger
-
-if TYPE_CHECKING:
-    import daft
-    import pandas as pd
-
-logger = get_logger(__name__)
-
-
-class JsonInput(Input):
-    """
-    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,
-        chunk_size: int = 100000,
-    ):
-        """Initialize the JsonInput class.
-
-        Args:
-            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
-        self.chunk_size = chunk_size
-        self.file_names = file_names
-
-    async def get_batched_dataframe(
-        self,
-    ) -> Union[AsyncIterator["pd.DataFrame"], Iterator["pd.DataFrame"]]:
-        """
-        Method to read the data from the json files in the path
-        and return as a batched pandas dataframe
-        """
-        try:
-            import pandas as pd
-
-            # 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 json_file in json_files:
-                json_reader_obj = pd.read_json(
-                    json_file,
-                    chunksize=self.chunk_size,
-                    lines=True,
-                )
-                for chunk in json_reader_obj:
-                    yield chunk
-        except Exception as e:
-            logger.error(f"Error reading batched data from JSON: {str(e)}")
-            raise
-
-    async def get_dataframe(self) -> "pd.DataFrame":
-        """
-        Method to read the data from the json files in the path
-        and return as a single combined pandas dataframe
-        """
-        try:
-            import pandas as pd
-
-            # 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(
-        self,
-    ) -> Union[AsyncIterator["daft.DataFrame"], Iterator["daft.DataFrame"]]:  # noqa: F821
-        """
-        Method to read the data from the json files in the path
-        and return as a batched daft dataframe
-        """
-        try:
-            import daft
-
-            # 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
-        """
-        Method to read the data from the json files in the path
-        and return as a single combined daft dataframe
-        """
-        try:
-            import daft
-
-            # 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,272 +0,0 @@
-from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Union
-
-from application_sdk.inputs import Input
-from application_sdk.observability.logger_adaptor import get_logger
-
-logger = get_logger(__name__)
-
-if TYPE_CHECKING:
-    import daft  # type: ignore
-    import pandas as pd
-
-
-class ParquetInput(Input):
-    """
-    Parquet Input class to read data from Parquet files using daft and pandas.
-    Supports reading both single files and directories containing multiple parquet files.
-    """
-
-    _EXTENSION = ".parquet"
-
-    def __init__(
-        self,
-        path: str,
-        chunk_size: int = 100000,
-        buffer_size: int = 5000,
-        file_names: Optional[List[str]] = None,
-    ):
-        """Initialize the Parquet input class.
-
-        Args:
-            path (str): Path to parquet file or directory containing parquet files.
-                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.
-            buffer_size (int): Number of rows per batch. Defaults to 5000.
-            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.buffer_size = buffer_size
-        self.file_names = file_names
-
-    async def get_dataframe(self) -> "pd.DataFrame":
-        """Read data from parquet file(s) and return as pandas DataFrame.
-
-        Returns:
-            pd.DataFrame: Combined dataframe from specified parquet files
-
-        Raises:
-            ValueError: When no valid path can be determined or no matching files found
-            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 files are read and combined
-        - Column schemas must be compatible across files
-        - Only reads files in the specified directory
-        """
-        try:
-            import pandas as pd
-
-            # 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)}")
-            raise
-
-    async def get_batched_dataframe(
-        self,
-    ) -> Union[AsyncIterator["pd.DataFrame"], Iterator["pd.DataFrame"]]:
-        """Read data from parquet file(s) in batches as pandas DataFrames.
-
-        Returns:
-            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
-
-            # 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]
-        except Exception as e:
-            logger.error(
-                f"Error reading data from parquet file(s) in batches: {str(e)}"
-            )
-            raise
-
-    async def get_daft_dataframe(self) -> "daft.DataFrame":  # noqa: F821
-        """Read data from parquet file(s) and return as daft DataFrame.
-
-        Returns:
-            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
-
-            # 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)}"
-            )
-            raise
-
-    async def get_batched_daft_dataframe(self) -> AsyncIterator["daft.DataFrame"]:  # type: ignore
-        """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 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
-
-            # 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")
-
-            # Create a lazy dataframe without loading data into memory
-            lazy_df = daft.read_parquet(parquet_files)
-
-            # Get total count efficiently
-            total_rows = lazy_df.count_rows()
-
-            # Yield chunks without loading everything into memory
-            for offset in range(0, total_rows, self.buffer_size):
-                chunk = lazy_df.offset(offset).limit(self.buffer_size)
-                yield chunk
-
-            del lazy_df
-
-        except Exception as error:
-            logger.error(
-                f"Error reading data from parquet file(s) in batches using daft: {error}"
-            )
-            raise