atlan-application-sdk 0.1.1rc36__tar.gz → 0.1.1rc37__tar.gz

atlan_application_sdk-0.1.1rc37/.github/release_notes.md

@@ -0,0 +1,7 @@
+## v0.1.1rc37 (September 03, 2025)
+
+Full Changelog: https://github.com/atlanhq/application-sdk/compare/v0.1.1rc36...v0.1.1rc37
+
+### Features
+
+- add file size-based chunking to JsonOutput (#650) (by @Nishchith Shetty in [7e7685b](https://github.com/atlanhq/application-sdk/commit/7e7685b))

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/CHANGELOG.md

@@ -1,5 +1,14 @@
 # Changelog
 
+## v0.1.1rc37 (September 03, 2025)
+
+Full Changelog: https://github.com/atlanhq/application-sdk/compare/v0.1.1rc36...v0.1.1rc37
+
+### Features
+
+- add file size-based chunking to JsonOutput (#650) (by @Nishchith Shetty in [7e7685b](https://github.com/atlanhq/application-sdk/commit/7e7685b))
+
+
 
 ## v0.1.1rc36 (September 01, 2025)
 

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 0.1.1rc36
+Version: 0.1.1rc37
 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.1rc36 → atlan_application_sdk-0.1.1rc37}/application_sdk/activities/common/models.py

@@ -4,7 +4,7 @@ This module contains Pydantic models used to represent various data structures
 needed by activities, such as statistics and configuration.
 """
 
-from typing import Optional
+from typing import List, Optional
 
 from pydantic import BaseModel
 
@@ -34,4 +34,5 @@ class ActivityStatistics(BaseModel):
 
     total_record_count: int = 0
     chunk_count: int = 0
+    partitions: Optional[List[int]] = []
     typename: Optional[str] = None

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/application_sdk/inputs/parquet.py

@@ -10,7 +10,7 @@ from application_sdk.services.objectstore import ObjectStore
 logger = get_logger(__name__)
 
 if TYPE_CHECKING:
-    import daft
+    import daft  # type: ignore
     import pandas as pd
 
 
@@ -136,7 +136,7 @@ class ParquetInput(Input):
             daft.DataFrame: Combined daft dataframe from all parquet files.
         """
         try:
-            import daft
+            import daft  # type: ignore
 
             if self.file_names:
                 path = f"{self.path}/{self.file_names[0].split('/')[0]}"
@@ -161,11 +161,11 @@ class ParquetInput(Input):
             a batch of data from the parquet file(s).
         """
         try:
-            import daft
+            import daft  # type: ignore
 
             if self.file_names:
                 for file_name in self.file_names:
-                    path = f"{self.path}/{file_name.replace('.json', '.parquet')}"
+                    path = f"{self.path}/{file_name}"
                     if self.input_prefix and path:
                         await self.download_files(path)
                         yield daft.read_parquet(path)

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/application_sdk/outputs/__init__.py

@@ -13,6 +13,7 @@ from typing import (
     Dict,
     Generator,
     List,
+    Literal,
     Optional,
     Union,
     cast,
@@ -31,7 +32,7 @@ logger = get_logger(__name__)
 activity.logger = logger
 
 if TYPE_CHECKING:
-    import daft
+    import daft  # type: ignore
     import pandas as pd
 
 
@@ -52,6 +53,27 @@ class Output(ABC):
     output_prefix: str
     total_record_count: int
     chunk_count: int
+    statistics: List[int] = []
+
+    def estimate_dataframe_file_size(
+        self, dataframe: "pd.DataFrame", file_type: Literal["json", "parquet"]
+    ) -> int:
+        """Estimate File size of a DataFrame by sampling a few records."""
+        if len(dataframe) == 0:
+            return 0
+
+        # Sample up to 10 records to estimate average size
+        sample_size = min(10, len(dataframe))
+        sample = dataframe.head(sample_size)
+        if file_type == "json":
+            sample_file = sample.to_json(orient="records", lines=True)
+        else:
+            sample_file = sample.to_parquet(index=False, compression="snappy")
+        if sample_file is not None:
+            avg_record_size = len(sample_file) / sample_size
+            return int(avg_record_size * len(dataframe))
+
+        return 0
 
     def process_null_fields(
         self,
@@ -217,6 +239,7 @@ class Output(ABC):
             statistics = {
                 "total_record_count": self.total_record_count,
                 "chunk_count": self.chunk_count,
+                "partitions": self.statistics,
             }
 
             # Write the statistics to a json file

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/application_sdk/outputs/json.py

@@ -6,6 +6,7 @@ import orjson
 from temporalio import activity
 
 from application_sdk.activities.common.utils import get_object_store_prefix
+from application_sdk.constants import DAPR_MAX_GRPC_MESSAGE_LENGTH
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.observability.metrics_adaptor import MetricType, get_metrics
 from application_sdk.outputs import Output
@@ -15,7 +16,7 @@ logger = get_logger(__name__)
 activity.logger = logger
 
 if TYPE_CHECKING:
-    import daft
+    import daft  # type: ignore
     import pandas as pd
 
 
@@ -32,7 +33,7 @@ def path_gen(chunk_start: int | None, chunk_count: int) -> str:
     if chunk_start is None:
         return f"{str(chunk_count)}.json"
     else:
-        return f"{str(chunk_start+chunk_count)}.json"
+        return f"chunk-{chunk_start}-part{chunk_count}.json"
 
 
 def convert_datetime_to_epoch(data: Any) -> Any:
@@ -124,6 +125,10 @@ class JsonOutput(Output):
         self.chunk_size = chunk_size or 100000
         self.buffer: List[Union["pd.DataFrame", "daft.DataFrame"]] = []  # noqa: F821
         self.current_buffer_size = 0
+        self.current_buffer_size_bytes = 0  # Track estimated buffer size in bytes
+        self.max_file_size_bytes = int(
+            DAPR_MAX_GRPC_MESSAGE_LENGTH * 0.9
+        )  # 90% of DAPR limit as safety buffer
         self.path_gen = path_gen
         self.start_marker = start_marker
         self.end_marker = end_marker
@@ -172,8 +177,21 @@ class JsonOutput(Output):
             ]
 
             for chunk in chunks:
+                # Estimate size of this chunk
+                chunk_size_bytes = self.estimate_dataframe_file_size(chunk, "json")
+
+                # Check if adding this chunk would exceed size limit
+                if (
+                    self.current_buffer_size_bytes + chunk_size_bytes
+                    > self.max_file_size_bytes
+                    and self.current_buffer_size > 0
+                ):
+                    # Flush current buffer before adding this chunk
+                    await self._flush_buffer()
+
                 self.buffer.append(chunk)
                 self.current_buffer_size += len(chunk)
+                self.current_buffer_size_bytes += chunk_size_bytes
 
                 if self.current_buffer_size >= partition:
                     await self._flush_buffer()
@@ -237,45 +255,19 @@ class JsonOutput(Output):
                     row, preserve_fields, null_to_empty_dict_fields
                 )
                 # Serialize the row and add it to the buffer
-                buffer.append(
-                    orjson.dumps(cleaned_row, option=orjson.OPT_APPEND_NEWLINE).decode(
-                        "utf-8"
-                    )
-                )
-
-            # If the buffer reaches the specified size, write it to the file
-            if self.chunk_size and len(buffer) >= self.chunk_size:
-                self.chunk_count += 1
-                output_file_name = f"{self.output_path}/{self.path_gen(self.chunk_start, self.chunk_count)}"
-                with open(output_file_name, "w") as f:
-                    f.writelines(buffer)
-                buffer.clear()  # Clear the buffer
-
-                # Record chunk metrics
-                self.metrics.record_metric(
-                    name="json_chunks_written",
-                    value=1,
-                    metric_type=MetricType.COUNTER,
-                    labels={"type": "daft"},
-                    description="Number of chunks written to JSON files",
-                )
+                serialized_row = orjson.dumps(
+                    cleaned_row, option=orjson.OPT_APPEND_NEWLINE
+                ).decode("utf-8")
+                buffer.append(serialized_row)
+                self.current_buffer_size_bytes += len(serialized_row)
+                if (self.chunk_size and len(buffer) >= self.chunk_size) or (
+                    self.current_buffer_size_bytes > self.max_file_size_bytes
+                ):
+                    await self.flush_daft_buffer(buffer)
 
             # Write any remaining rows in the buffer
             if buffer:
-                self.chunk_count += 1
-                output_file_name = f"{self.output_path}/{self.path_gen(self.chunk_start, self.chunk_count)}"
-                with open(output_file_name, "w") as f:
-                    f.writelines(buffer)
-                buffer.clear()
-
-                # Record chunk metrics
-                self.metrics.record_metric(
-                    name="json_chunks_written",
-                    value=1,
-                    metric_type=MetricType.COUNTER,
-                    labels={"type": "daft"},
-                    description="Number of chunks written to JSON files",
-                )
+                await self.flush_daft_buffer(buffer)
 
             # Record metrics for successful write
             self.metrics.record_metric(
@@ -303,6 +295,32 @@ class JsonOutput(Output):
             )
             logger.error(f"Error writing daft dataframe to json: {str(e)}")
 
+    async def flush_daft_buffer(self, buffer: List[str]):
+        """Flush the current buffer to a JSON file.
+
+        This method combines all DataFrames in the buffer, writes them to a JSON file,
+        and uploads the file to the object store.
+        """
+        self.chunk_count += 1
+        output_file_name = (
+            f"{self.output_path}/{self.path_gen(self.chunk_start, self.chunk_count)}"
+        )
+        with open(output_file_name, "w") as f:
+            f.writelines(buffer)
+        buffer.clear()  # Clear the buffer
+
+        self.current_buffer_size = 0
+        self.current_buffer_size_bytes = 0
+
+        # Record chunk metrics
+        self.metrics.record_metric(
+            name="json_chunks_written",
+            value=1,
+            metric_type=MetricType.COUNTER,
+            labels={"type": "daft"},
+            description="Number of chunks written to JSON files",
+        )
+
     async def _flush_buffer(self):
         """Flush the current buffer to a JSON file.
 
@@ -353,6 +371,7 @@ class JsonOutput(Output):
 
             self.buffer.clear()
             self.current_buffer_size = 0
+            self.current_buffer_size_bytes = 0
 
         except Exception as e:
             # Record metrics for failed write

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/application_sdk/outputs/parquet.py

@@ -1,9 +1,10 @@
 import os
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import TYPE_CHECKING, List, Literal, Optional, Union
 
 from temporalio import activity
 
 from application_sdk.activities.common.utils import get_object_store_prefix
+from application_sdk.constants import DAPR_MAX_GRPC_MESSAGE_LENGTH
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.observability.metrics_adaptor import MetricType, get_metrics
 from application_sdk.outputs import Output
@@ -13,7 +14,7 @@ logger = get_logger(__name__)
 activity.logger = logger
 
 if TYPE_CHECKING:
-    import daft
+    import daft  # type: ignore
     import pandas as pd
 
 
@@ -46,6 +47,7 @@ class ParquetOutput(Output):
         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,
         chunk_count: int = 0,
         chunk_start: Optional[int] = None,
@@ -78,11 +80,19 @@ class ParquetOutput(Output):
         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
         self.total_record_count = total_record_count
         self.chunk_count = chunk_count
+        self.current_buffer_size = 0
+        self.current_buffer_size_bytes = 0  # Track estimated buffer size in bytes
+        self.max_file_size_bytes = int(
+            DAPR_MAX_GRPC_MESSAGE_LENGTH * 0.9
+        )  # 90% of DAPR limit as safety buffer
         self.chunk_start = chunk_start
         self.start_marker = start_marker
         self.end_marker = end_marker
+        self.statistics = []
         self.metrics = get_metrics()
 
         # Create output directory
@@ -117,7 +127,7 @@ class ParquetOutput(Output):
         if chunk_start is None:
             return f"{str(chunk_count)}.parquet"
         else:
-            return f"{str(chunk_start+chunk_count)}.parquet"
+            return f"chunk-{str(chunk_start)}-part{str(chunk_count)}.parquet"
 
     async def write_dataframe(self, dataframe: "pd.DataFrame"):
         """Write a pandas DataFrame to Parquet files and upload to object store.
@@ -126,20 +136,46 @@ class ParquetOutput(Output):
             dataframe (pd.DataFrame): The DataFrame to write.
         """
         try:
+            chunk_part = 0
             if len(dataframe) == 0:
                 return
 
-            # Update counters
-            self.chunk_count += 1
-            self.total_record_count += len(dataframe)
-            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 pandas native method
-            dataframe.to_parquet(
-                file_path,
-                index=False,
-                compression="snappy",  # Using snappy compression by default
+            # Split the DataFrame into chunks
+            partition = (
+                self.chunk_size
+                if self.chunk_start is None
+                else min(self.chunk_size, self.buffer_size)
             )
+            chunks = [
+                dataframe[i : i + partition]  # type: ignore
+                for i in range(0, len(dataframe), partition)
+            ]
+
+            for chunk in chunks:
+                # Estimate size of this chunk
+                chunk_size_bytes = self.estimate_dataframe_file_size(chunk, "parquet")
+
+                # Check if adding this chunk would exceed size limit
+                if (
+                    self.current_buffer_size_bytes + chunk_size_bytes
+                    > self.max_file_size_bytes
+                    and self.current_buffer_size > 0
+                ):
+                    # Flush current buffer before adding this chunk
+                    chunk_part += 1
+                    await self._flush_buffer(chunk_part)
+
+                self.buffer.append(chunk)
+                self.current_buffer_size += len(chunk)
+                self.current_buffer_size_bytes += chunk_size_bytes
+
+                if self.current_buffer_size >= partition:  # type: ignore
+                    chunk_part += 1
+                    await self._flush_buffer(chunk_part)
+
+            if self.buffer and self.current_buffer_size > 0:
+                chunk_part += 1
+                await self._flush_buffer(chunk_part)
 
             # Record metrics for successful write
             self.metrics.record_metric(
@@ -159,11 +195,8 @@ class ParquetOutput(Output):
                 description="Number of chunks written to Parquet files",
             )
 
-            # Upload the file to object store
-            await ObjectStore.upload_file(
-                source=file_path,
-                destination=get_object_store_prefix(file_path),
-            )
+            self.chunk_count += 1
+            self.statistics.append(chunk_part)
         except Exception as e:
             # Record metrics for failed write
             self.metrics.record_metric(
@@ -245,3 +278,68 @@ class ParquetOutput(Output):
             str: The full path of the output file.
         """
         return self.output_path
+
+    async def _flush_buffer(self, chunk_part):
+        """Flush the current buffer to a Parquet file.
+
+        This method combines all DataFrames in the buffer, writes them to a Parquet file,
+        and uploads the file to the object store.
+
+        Note:
+            If the buffer is empty or has no records, the method returns without writing.
+        """
+        import pandas as pd
+
+        if not self.buffer or not self.current_buffer_size:
+            return
+
+        if not all(isinstance(df, pd.DataFrame) for df in self.buffer):
+            raise TypeError(
+                "_flush_buffer encountered non-DataFrame elements in buffer. This should not happen."
+            )
+
+        try:
+            # Now it's safe to cast for pd.concat
+            pd_buffer: List[pd.DataFrame] = self.buffer  # type: ignore
+            combined_dataframe = pd.concat(pd_buffer)
+
+            # Write DataFrame to Parquet file
+            if not combined_dataframe.empty:
+                self.total_record_count += len(combined_dataframe)
+                output_file_name = (
+                    f"{self.output_path}/{self.path_gen(self.chunk_count, chunk_part)}"
+                )
+                combined_dataframe.to_parquet(
+                    output_file_name, index=False, compression="snappy"
+                )
+
+                # Record chunk metrics
+                self.metrics.record_metric(
+                    name="parquet_chunks_written",
+                    value=1,
+                    metric_type=MetricType.COUNTER,
+                    labels={"type": "pandas"},
+                    description="Number of chunks written to Parquet files",
+                )
+
+                # Push the file to the object store
+                await ObjectStore.upload_file(
+                    source=output_file_name,
+                    destination=get_object_store_prefix(output_file_name),
+                )
+
+            self.buffer.clear()
+            self.current_buffer_size = 0
+            self.current_buffer_size_bytes = 0
+
+        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": "pandas", "error": str(e)},
+                description="Number of errors while writing to Parquet files",
+            )
+            logger.error(f"Error flushing buffer to parquet: {str(e)}")
+            raise e

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/application_sdk/version.py

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

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/application_sdk/workflows/metadata_extraction/sql.py

@@ -107,7 +107,11 @@ class BaseSQLMetadataExtractionWorkflow(MetadataExtractionWorkflow):
         activity_statistics = ActivityStatistics.model_validate(raw_statistics)
         transform_activities: List[Any] = []
 
-        if activity_statistics is None or activity_statistics.chunk_count == 0:
+        if (
+            activity_statistics is None
+            or activity_statistics.chunk_count == 0
+            or not activity_statistics.partitions
+        ):
             # to handle the case where the fetch_fn returns None or no chunks
             return
 
@@ -115,7 +119,9 @@ class BaseSQLMetadataExtractionWorkflow(MetadataExtractionWorkflow):
             raise ValueError("Invalid typename")
 
         batches, chunk_starts = self.get_transform_batches(
-            activity_statistics.chunk_count, activity_statistics.typename
+            activity_statistics.chunk_count,
+            activity_statistics.typename,
+            activity_statistics.partitions,
         )
 
         for i in range(len(batches)):
@@ -144,7 +150,9 @@ class BaseSQLMetadataExtractionWorkflow(MetadataExtractionWorkflow):
             total_record_count += metadata_model.total_record_count
             chunk_count += metadata_model.chunk_count
 
-    def get_transform_batches(self, chunk_count: int, typename: str):
+    def get_transform_batches(
+        self, chunk_count: int, typename: str, partitions: List[int]
+    ):
         """Get batches for parallel transformation processing.
 
         Args:
@@ -159,12 +167,17 @@ class BaseSQLMetadataExtractionWorkflow(MetadataExtractionWorkflow):
         batches: List[List[str]] = []
         chunk_start_numbers: List[int] = []
 
-        for i in range(chunk_count):
+        for i, partition in enumerate(partitions):
             # Track starting chunk number (which is just i)
             chunk_start_numbers.append(i)
 
             # Each batch contains exactly one chunk
-            batches.append([f"{typename}/{i+1}.json"])
+            batches.append(
+                [
+                    f"{typename}/chunk-{i}-part{file+1}.parquet"
+                    for file in range(partition)
+                ]
+            )
 
         return batches, chunk_start_numbers
 

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "atlan-application-sdk"
-version = "0.1.1rc36"
+version = "0.1.1rc37"
 description = "Atlan Application SDK is a Python library for developing applications on the Atlan Platform"
 license = "Apache-2.0"
 authors = [

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/tests/unit/inputs/test_parquet.py

@@ -620,7 +620,10 @@ async def test_get_batched_daft_dataframe_with_file_names(monkeypatch) -> None:
         "glob.glob", return_value=[]
     ), patch("application_sdk.services.objectstore.ObjectStore.download_file"):
         path = "/data"
-        file_names = ["one.json", "two.json"]  # Note: .json extension gets replaced
+        file_names = [
+            "one.parquet",
+            "two.parquet",
+        ]  # Note: .json extension gets replaced
         input_prefix = "remote"
 
         parquet_input = ParquetInput(

{atlan_application_sdk-0.1.1rc36 → atlan_application_sdk-0.1.1rc37}/tests/unit/outputs/test_output.py

@@ -107,11 +107,12 @@ class TestOutput:
         """Test write_statistics successful case."""
         self.output.total_record_count = 100
         self.output.chunk_count = 5
+        self.output.statistics = [1, 2, 1, 2, 1]
 
         # Mock the open function, orjson.dumps, and object store upload
         with patch("builtins.open", mock_open()) as mock_file, patch(
             "orjson.dumps",
-            return_value=b'{"total_record_count": 100, "chunk_count": 5}',
+            return_value=b'{"total_record_count": 100, "chunk_count": 5, "partitions": [1,2,1,2,1]}',
         ) as mock_orjson, patch(
             "application_sdk.outputs.get_object_store_prefix",
             return_value="path/statistics.json.ignore",
@@ -123,10 +124,18 @@ class TestOutput:
             stats = await self.output.write_statistics()
 
             # Assertions
-            assert stats == {"total_record_count": 100, "chunk_count": 5}
+            assert stats == {
+                "total_record_count": 100,
+                "chunk_count": 5,
+                "partitions": [1, 2, 1, 2, 1],
+            }
             mock_file.assert_called_once_with("/test/path/statistics.json.ignore", "w")
             mock_orjson.assert_called_once_with(
-                {"total_record_count": 100, "chunk_count": 5}
+                {
+                    "total_record_count": 100,
+                    "chunk_count": 5,
+                    "partitions": [1, 2, 1, 2, 1],
+                }
             )
             # Verify the upload call
             mock_push.assert_awaited_once()