atlan-application-sdk 0.1.1rc38__py3-none-any.whl → 0.1.1rc40__py3-none-any.whl

application_sdk/outputs/.cursor/BUGBOT.md

@@ -0,0 +1,295 @@
+# Output Code Review Guidelines - Data Output Processing
+
+## Context-Specific Patterns
+
+This directory contains output processing implementations for various data formats (JSON, Parquet, Iceberg). Output processors must handle data uploads efficiently while maintaining data integrity and correct destination paths.
+
+### Phase 1: Critical Output Safety Issues
+
+**Object Store Path Management:**
+
+- **Correct destination paths**: Upload paths must respect user-configured output prefixes
+- **Path construction accuracy**: Object store keys must be calculated correctly, not hardcoded
+- **User prefix preservation**: Respect user-provided output directories and naming conventions
+- **Path validation**: Ensure upload paths don't conflict with existing data
+
+**Data Integrity and Security:**
+
+- All output data must be validated before upload
+- File permissions and access controls must be properly set
+- Data serialization must be consistent and recoverable
+- Prevent overwriting critical data without confirmation
+- Maintain data lineage information in output metadata
+
+```python
+# ✅ DO: Proper object store upload path handling
+class JsonOutput:
+    async def upload_to_object_store(
+        self,
+        data: List[dict],
+        output_prefix: str,  # User-provided output location
+        filename: str
+    ) -> dict:
+        """Upload data with correct path handling."""
+
+        # Construct full object store path respecting user's output prefix
+        object_store_key = os.path.join(output_prefix, filename)
+
+        # Serialize data
+        json_data = orjson.dumps(data, option=orjson.OPT_APPEND_NEWLINE)
+
+        # Upload to correct location
+        result = await self.object_store.upload_file(
+            data=json_data,
+            destination=object_store_key  # Respect user's intended location
+        )
+
+        return result
+
+# ❌ REJECT: Incorrect path handling
+class BadJsonOutput:
+    async def upload_to_object_store(self, data: List[dict], filename: str):
+        # Wrong: hardcoded or derived path, ignoring user configuration
+        object_store_key = get_object_store_prefix(f"/tmp/{filename}")  # Ignores output_prefix!
+
+        result = await self.object_store.upload_file(
+            data=orjson.dumps(data),
+            destination=object_store_key  # Wrong destination!
+        )
+        return result
+```
+
+### Phase 2: Output Architecture Patterns
+
+**Performance Optimization Requirements:**
+
+- **Parallelization opportunities**: Flag sequential upload operations that could be parallelized
+- **Batch processing**: Group related uploads to reduce overhead
+- **Streaming uploads**: Use streaming for large datasets instead of loading into memory
+- **Connection optimization**: Reuse object store connections across operations
+
+**Resource Management:**
+
+- Use proper connection pooling for object store operations
+- Implement timeout handling for upload operations
+- Clean up temporary files after upload
+- Handle partial upload failures gracefully
+- Monitor memory usage during large data serialization
+
+```python
+# ✅ DO: Parallel upload processing
+async def upload_multiple_datasets_parallel(
+    self,
+    datasets: List[Tuple[List[dict], str]],  # (data, filename) pairs
+    output_prefix: str
+) -> List[dict]:
+    """Upload multiple datasets in parallel for better performance."""
+
+    async def upload_single_dataset(data: List[dict], filename: str) -> dict:
+        """Upload a single dataset with error handling."""
+        try:
+            object_store_key = os.path.join(output_prefix, filename)
+            serialized_data = orjson.dumps(data, option=orjson.OPT_APPEND_NEWLINE)
+
+            return await self.object_store.upload_file(
+                data=serialized_data,
+                destination=object_store_key
+            )
+        except Exception as e:
+            logger.error(f"Failed to upload {filename}: {e}")
+            raise
+
+    # Parallel processing with controlled concurrency
+    semaphore = asyncio.Semaphore(5)  # Limit concurrent uploads
+
+    async def upload_with_semaphore(data: List[dict], filename: str) -> dict:
+        async with semaphore:
+            return await upload_single_dataset(data, filename)
+
+    tasks = [upload_with_semaphore(data, filename) for data, filename in datasets]
+    return await asyncio.gather(*tasks)
+
+# ❌ REJECT: Sequential upload processing
+async def upload_multiple_datasets_sequential(
+    self,
+    datasets: List[Tuple[List[dict], str]],
+    output_prefix: str
+) -> List[dict]:
+    """Sequential uploads - should be flagged for parallelization."""
+    results = []
+    for data, filename in datasets:  # FLAG: Could be parallelized
+        object_store_key = os.path.join(output_prefix, filename)
+        result = await self.object_store.upload_file(data, object_store_key)
+        results.append(result)
+    return results
+```
+
+### Phase 3: Output Testing Requirements
+
+**Data Output Testing:**
+
+- Test with various data formats and sizes
+- Test serialization and deserialization consistency
+- Test partial upload scenarios and recovery
+- Mock object store operations in unit tests
+- Include integration tests with real object store
+- Test data corruption detection and prevention
+
+**Performance Testing:**
+
+- Include tests for large dataset uploads
+- Test memory usage during serialization
+- Test concurrent upload operations
+- Verify timeout handling works correctly
+- Test connection pool behavior under load
+
+### Phase 4: Performance and Scalability
+
+**Data Upload Efficiency:**
+
+- Use streaming uploads for large datasets
+- Implement proper chunking for oversized data
+- Use compression for large text-based outputs
+- Monitor upload progress and provide feedback
+- Optimize serialization performance (use orjson over json)
+
+**Object Store Optimization:**
+
+- Use connection pooling for object store clients
+- Implement proper retry logic for upload failures
+- Use parallel uploads where appropriate
+- Monitor upload metrics and error rates
+- Handle bandwidth limitations gracefully
+
+### Phase 5: Output Maintainability
+
+**Error Handling and Recovery:**
+
+- Implement comprehensive error handling for all upload operations
+- Provide meaningful error messages with upload context
+- Handle partial upload failures gracefully
+- Implement proper retry logic for transient failures
+- Log all upload operations with destination information
+
+**Configuration Management:**
+
+- Externalize all output-related configuration
+- Support different output destinations and formats
+- Validate output configuration before processing
+- Document all supported output parameters
+- Handle environment-specific output requirements
+
+---
+
+## Output-Specific Anti-Patterns
+
+**Always Reject:**
+
+- **Path derivation errors**: Deriving object store paths from local temporary paths
+- **Sequential uploads**: Uploading multiple files sequentially when parallel uploads are possible
+- **Memory inefficiency**: Loading entire datasets into memory for serialization
+- **Missing upload verification**: Not verifying successful uploads
+- **Poor error recovery**: Not handling partial upload failures gracefully
+- **Resource leaks**: Not cleaning up temporary files or connections
+
+**Object Store Upload Anti-Patterns:**
+
+```python
+# ❌ REJECT: Incorrect upload path handling
+class BadOutputProcessor:
+    async def upload_results(self, results: List[dict]):
+        # Wrong: derives upload path from temporary local path
+        local_temp_file = "/tmp/results.json"
+        upload_key = get_object_store_prefix(local_temp_file)  # Incorrect!
+
+        await self.object_store.upload_file(results, upload_key)
+
+# ✅ REQUIRE: Correct upload path handling
+class GoodOutputProcessor:
+    async def upload_results(
+        self,
+        results: List[dict],
+        output_prefix: str,  # User-specified destination
+        filename: str = "results.json"
+    ):
+        # Use actual user-configured output location
+        upload_key = os.path.join(output_prefix, filename)
+
+        await self.object_store.upload_file(
+            data=orjson.dumps(results),
+            destination=upload_key  # Correct destination
+        )
+```
+
+**Performance Anti-Patterns:**
+
+```python
+# ❌ REJECT: Sequential upload processing
+async def upload_multiple_files_sequential(file_data_pairs: List[Tuple]):
+    results = []
+    for data, filename in file_data_pairs:  # Should be parallelized
+        result = await upload_single_file(data, filename)
+        results.append(result)
+    return results
+
+# ✅ REQUIRE: Parallel upload processing with proper error handling
+async def upload_multiple_files_parallel(
+    file_data_pairs: List[Tuple],
+    max_concurrency: int = 5
+) -> List[dict]:
+    semaphore = asyncio.Semaphore(max_concurrency)
+
+    async def upload_with_semaphore(data, filename):
+        async with semaphore:
+            try:
+                return await upload_single_file(data, filename)
+            except Exception as e:
+                logger.error(f"Upload failed for {filename}: {e}")
+                return {"filename": filename, "status": "failed", "error": str(e)}
+
+    tasks = [upload_with_semaphore(data, filename) for data, filename in file_data_pairs]
+    return await asyncio.gather(*tasks)
+```
+
+**Memory Management Anti-Patterns:**
+
+```python
+# ❌ REJECT: Loading entire dataset for serialization
+async def bad_large_dataset_upload(large_dataset: List[dict]):
+    # Loads entire dataset into memory
+    json_data = orjson.dumps(large_dataset)  # Could exceed memory limits
+    await upload_data(json_data)
+
+# ✅ REQUIRE: Streaming serialization for large datasets
+async def good_large_dataset_upload(large_dataset: List[dict], chunk_size: int = 1000):
+    """Stream large datasets to avoid memory issues."""
+
+    async def serialize_chunk(chunk: List[dict]) -> bytes:
+        return orjson.dumps(chunk, option=orjson.OPT_APPEND_NEWLINE)
+
+    # Process in chunks to manage memory
+    for i in range(0, len(large_dataset), chunk_size):
+        chunk = large_dataset[i:i + chunk_size]
+        serialized_chunk = await serialize_chunk(chunk)
+
+        await upload_chunk(
+            data=serialized_chunk,
+            chunk_index=i // chunk_size
+        )
+```
+
+## Educational Context for Output Reviews
+
+When reviewing output code, emphasize:
+
+1. **Data Integrity Impact**: "Incorrect upload path handling can cause data to be stored in wrong locations, making it inaccessible to downstream processes. This breaks data pipelines and can cause data loss."
+
+2. **Performance Impact**: "Sequential uploads create unnecessary bottlenecks. For enterprise datasets with multiple output files, parallelization can significantly reduce processing time and improve user experience."
+
+3. **Resource Impact**: "Poor memory management during serialization can cause out-of-memory errors with large datasets. Streaming and chunking are essential for enterprise-scale data output."
+
+4. **User Experience Impact**: "Output path errors are often discovered late in processing, causing wasted computation and frustrating delays. Proper validation and clear error messages improve reliability."
+
+5. **Scalability Impact**: "Output patterns that work for small datasets can fail at enterprise scale. Always design output processes to handle the largest expected dataset sizes efficiently."
+
+6. **Data Pipeline Impact**: "Output processing is the final step in data pipelines. Failures here can invalidate all upstream processing work. Robust error handling and verification are critical for pipeline reliability."

application_sdk/outputs/iceberg.py

@@ -29,6 +29,7 @@ class IcebergOutput(Output):
         mode: str = "append",
         total_record_count: int = 0,
         chunk_count: int = 0,
+        retain_local_copy: bool = False,
     ):
         """Initialize the Iceberg output class.
 
@@ -39,6 +40,8 @@ class IcebergOutput(Output):
             mode (str, optional): Write mode for the iceberg table. Defaults to "append".
             total_record_count (int, optional): Total record count written to the iceberg table. Defaults to 0.
             chunk_count (int, optional): Number of chunks written to the iceberg table. Defaults to 0.
+            retain_local_copy (bool, optional): Whether to retain the local copy of the files.
+                Defaults to False.
         """
         self.total_record_count = total_record_count
         self.chunk_count = chunk_count
@@ -47,6 +50,7 @@ class IcebergOutput(Output):
         self.iceberg_table = iceberg_table
         self.mode = mode
         self.metrics = get_metrics()
+        self.retain_local_copy = retain_local_copy
 
     async def write_dataframe(self, dataframe: "pd.DataFrame"):
         """

application_sdk/outputs/json.py

@@ -93,6 +93,7 @@ class JsonOutput(Output):
         path_gen: Callable[[int | None, int], str] = path_gen,
         start_marker: Optional[str] = None,
         end_marker: Optional[str] = None,
+        retain_local_copy: bool = False,
         **kwargs: Dict[str, Any],
     ):
         """Initialize the JSON output handler.
@@ -113,6 +114,8 @@ class JsonOutput(Output):
                 Defaults to 0.
             path_gen (Callable, optional): Function to generate file paths.
                 Defaults to path_gen function.
+            retain_local_copy (bool, optional): Whether to retain the local copy of the files.
+                Defaults to False.
         """
         self.output_path = output_path
         self.output_suffix = output_suffix
@@ -133,6 +136,7 @@ class JsonOutput(Output):
         self.start_marker = start_marker
         self.end_marker = end_marker
         self.metrics = get_metrics()
+        self.retain_local_copy = retain_local_copy
 
         if not self.output_path:
             raise ValueError("output_path is required")
@@ -282,6 +286,7 @@ class JsonOutput(Output):
             await ObjectStore.upload_prefix(
                 source=self.output_path,
                 destination=get_object_store_prefix(self.output_path),
+                retain_local_copy=self.retain_local_copy,
             )
 
         except Exception as e:
@@ -367,6 +372,7 @@ class JsonOutput(Output):
                 await ObjectStore.upload_file(
                     source=output_file_name,
                     destination=get_object_store_prefix(output_file_name),
+                    retain_local_copy=self.retain_local_copy,
                 )
 
             self.buffer.clear()

application_sdk/outputs/parquet.py

@@ -1,5 +1,6 @@
 import os
-from typing import TYPE_CHECKING, List, Literal, Optional, Union
+from enum import Enum
+from typing import TYPE_CHECKING, List, Optional, Union
 
 from temporalio import activity
 
@@ -18,6 +19,14 @@ if TYPE_CHECKING:
     import pandas as pd
 
 
+class WriteMode(Enum):
+    """Enumeration of write modes for Parquet output operations."""
+
+    APPEND = "append"
+    OVERWRITE = "overwrite"
+    OVERWRITE_PARTITIONS = "overwrite-partitions"
+
+
 class ParquetOutput(Output):
     """Output handler for writing data to Parquet files.
 
@@ -29,7 +38,6 @@ class ParquetOutput(Output):
         output_prefix (str): Prefix for files when uploading to object store.
         output_suffix (str): Suffix for output files.
         typename (Optional[str]): Type name of the entity e.g database, schema, table.
-        mode (str): Write mode for parquet files ("append" or "overwrite").
         chunk_size (int): Maximum number of records per chunk.
         total_record_count (int): Total number of records processed.
         chunk_count (int): Number of chunks created.
@@ -45,7 +53,6 @@ class ParquetOutput(Output):
         output_suffix: str = "",
         output_prefix: str = "",
         typename: Optional[str] = None,
-        write_mode: Literal["append", "overwrite", "overwrite-partitions"] = "append",
         chunk_size: Optional[int] = 100000,
         buffer_size: Optional[int] = 100000,
         total_record_count: int = 0,
@@ -53,6 +60,7 @@ class ParquetOutput(Output):
         chunk_start: Optional[int] = None,
         start_marker: Optional[str] = None,
         end_marker: Optional[str] = None,
+        retain_local_copy: bool = False,
     ):
         """Initialize the Parquet output handler.
 
@@ -61,7 +69,6 @@ class ParquetOutput(Output):
             output_suffix (str): Suffix for output files.
             output_prefix (str): Prefix for files when uploading to object store.
             typename (Optional[str], optional): Type name of the entity e.g database, schema, table.
-            mode (str, optional): Write mode for parquet files. Defaults to "append".
             chunk_size (int, optional): Maximum records per chunk. Defaults to 100000.
             total_record_count (int, optional): Initial total record count. Defaults to 0.
             chunk_count (int, optional): Initial chunk count. Defaults to 0.
@@ -73,12 +80,13 @@ class ParquetOutput(Output):
                 Defaults to None.
             end_marker (Optional[str], optional): End marker for query extraction.
                 Defaults to None.
+            retain_local_copy (bool, optional): Whether to retain the local copy of the files.
+                Defaults to False.
         """
         self.output_path = output_path
         self.output_suffix = output_suffix
         self.output_prefix = output_prefix
         self.typename = typename
-        self.write_mode = write_mode
         self.chunk_size = chunk_size
         self.buffer_size = buffer_size
         self.buffer: List[Union["pd.DataFrame", "daft.DataFrame"]] = []  # noqa: F821
@@ -94,6 +102,7 @@ class ParquetOutput(Output):
         self.end_marker = end_marker
         self.statistics = []
         self.metrics = get_metrics()
+        self.retain_local_copy = retain_local_copy
 
         # Create output directory
         self.output_path = os.path.join(self.output_path, self.output_suffix)
@@ -103,7 +112,7 @@ class ParquetOutput(Output):
 
     def path_gen(
         self,
-        chunk_start: int | None = None,
+        chunk_start: Optional[int] = None,
         chunk_count: int = 0,
         start_marker: Optional[str] = None,
         end_marker: Optional[str] = None,
@@ -111,7 +120,7 @@ class ParquetOutput(Output):
         """Generate a file path for a chunk.
 
         Args:
-            chunk_start (int | None): Starting index of the chunk, or None for single chunk.
+            chunk_start (Optional[int]): Starting index of the chunk, or None for single chunk.
             chunk_count (int): Total number of chunks.
             start_marker (Optional[str]): Start marker for query extraction.
             end_marker (Optional[str]): End marker for query extraction.
@@ -182,7 +191,7 @@ class ParquetOutput(Output):
                 name="parquet_write_records",
                 value=len(dataframe),
                 metric_type=MetricType.COUNTER,
-                labels={"type": "pandas", "mode": self.write_mode},
+                labels={"type": "pandas", "mode": WriteMode.APPEND.value},
                 description="Number of records written to Parquet files from pandas DataFrame",
             )
 
@@ -191,7 +200,7 @@ class ParquetOutput(Output):
                 name="parquet_chunks_written",
                 value=1,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "pandas", "mode": self.write_mode},
+                labels={"type": "pandas", "mode": WriteMode.APPEND.value},
                 description="Number of chunks written to Parquet files",
             )
 
@@ -203,69 +212,115 @@ class ParquetOutput(Output):
                 name="parquet_write_errors",
                 value=1,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "pandas", "mode": self.write_mode, "error": str(e)},
+                labels={
+                    "type": "pandas",
+                    "mode": WriteMode.APPEND.value,
+                    "error": str(e),
+                },
                 description="Number of errors while writing to Parquet files",
             )
             logger.error(f"Error writing pandas dataframe to parquet: {str(e)}")
             raise
 
-    async def write_daft_dataframe(self, dataframe: "daft.DataFrame"):  # noqa: F821
+    async def write_daft_dataframe(
+        self,
+        dataframe: "daft.DataFrame",  # noqa: F821
+        partition_cols: Optional[List] = None,
+        write_mode: Union[WriteMode, str] = WriteMode.APPEND,
+        morsel_size: int = 100_000,
+    ):
         """Write a daft DataFrame to Parquet files and upload to object store.
 
+        Uses Daft's native file size management to automatically split large DataFrames
+        into multiple parquet files based on the configured target file size. Supports
+        Hive partitioning for efficient data organization.
+
         Args:
             dataframe (daft.DataFrame): The DataFrame to write.
+            partition_cols (Optional[List]): Column names or expressions to use for Hive partitioning.
+                Can be strings (column names) or daft column expressions. If None (default), no partitioning is applied.
+            write_mode (Union[WriteMode, str]): Write mode for parquet files.
+                Use WriteMode.APPEND, WriteMode.OVERWRITE, WriteMode.OVERWRITE_PARTITIONS, or their string equivalents.
+            morsel_size (int): Default number of rows in a morsel used for the new local executor, when running locally on just a single machine,
+                Daft does not use partitions. Instead of using partitioning to control parallelism, the local execution engine performs a streaming-based
+                execution on small "morsels" of data, which provides much more stable memory utilization while improving the user experience with not having
+                to worry about partitioning.
+
+        Note:
+            - Daft automatically handles file chunking based on parquet_target_filesize
+            - Multiple files will be created if DataFrame exceeds DAPR limit
+            - If partition_cols is set, creates Hive-style directory structure
         """
         try:
+            import daft
+
+            # Convert string to enum if needed for backward compatibility
+            if isinstance(write_mode, str):
+                write_mode = WriteMode(write_mode)
+
             row_count = dataframe.count_rows()
             if row_count == 0:
                 return
 
+            # Use Daft's execution context for temporary configuration
+            with daft.execution_config_ctx(
+                parquet_target_filesize=self.max_file_size_bytes,
+                default_morsel_size=morsel_size,
+            ):
+                # Daft automatically handles file splitting and naming
+                dataframe.write_parquet(
+                    root_dir=self.output_path,
+                    write_mode=write_mode.value,
+                    partition_cols=partition_cols if partition_cols else [],
+                )
+
             # Update counters
             self.chunk_count += 1
             self.total_record_count += row_count
 
-            # Generate file path using path_gen function
-            if self.start_marker and self.end_marker:
-                file_path = self.output_path
-            else:
-                file_path = f"{self.output_path}/{self.path_gen(self.chunk_start, self.chunk_count, self.start_marker, self.end_marker)}"
-
-            # Write the dataframe to parquet using daft
-            dataframe.write_parquet(
-                file_path,
-                write_mode=self.write_mode,
-            )
-
             # Record metrics for successful write
             self.metrics.record_metric(
                 name="parquet_write_records",
                 value=row_count,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "daft", "mode": self.write_mode},
+                labels={"type": "daft", "mode": write_mode.value},
                 description="Number of records written to Parquet files from daft DataFrame",
             )
 
-            # Record chunk metrics
+            # Record operation metrics (note: actual file count may be higher due to Daft's splitting)
             self.metrics.record_metric(
-                name="parquet_chunks_written",
+                name="parquet_write_operations",
                 value=1,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "daft", "mode": self.write_mode},
-                description="Number of chunks written to Parquet files",
+                labels={"type": "daft", "mode": write_mode.value},
+                description="Number of write operations to Parquet files",
             )
 
-            # Upload the file to object store
-            await ObjectStore.upload_file(
-                source=file_path,
-                destination=get_object_store_prefix(file_path),
+            #  Upload the entire directory (contains multiple parquet files created by Daft)
+            if write_mode == WriteMode.OVERWRITE:
+                # Delete the directory from object store
+                try:
+                    await ObjectStore.delete_prefix(
+                        prefix=get_object_store_prefix(self.output_path)
+                    )
+                except FileNotFoundError as e:
+                    logger.info(
+                        f"No files found under prefix {get_object_store_prefix(self.output_path)}: {str(e)}"
+                    )
+
+            await ObjectStore.upload_prefix(
+                source=self.output_path,
+                destination=get_object_store_prefix(self.output_path),
+                retain_local_copy=self.retain_local_copy,
             )
+
         except Exception as e:
             # Record metrics for failed write
             self.metrics.record_metric(
                 name="parquet_write_errors",
                 value=1,
                 metric_type=MetricType.COUNTER,
-                labels={"type": "daft", "mode": self.write_mode, "error": str(e)},
+                labels={"type": "daft", "mode": write_mode, "error": str(e)},
                 description="Number of errors while writing to Parquet files",
             )
             logger.error(f"Error writing daft dataframe to parquet: {str(e)}")
@@ -279,7 +334,7 @@ class ParquetOutput(Output):
         """
         return self.output_path
 
-    async def _flush_buffer(self, chunk_part):
+    async def _flush_buffer(self, chunk_part: int):
         """Flush the current buffer to a Parquet file.
 
         This method combines all DataFrames in the buffer, writes them to a Parquet file,