data-factory-utils 3.2.2__tar.gz → 3.3.0__tar.gz

{data_factory_utils-3.2.2 → data_factory_utils-3.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: data-factory-utils
-Version: 3.2.2
+Version: 3.3.0
 Summary: Utility functions for interacting with data factories.
 Requires-Dist: boto3>=1.42.8
 Requires-Dist: boto3-stubs>=1.42.89

{data_factory_utils-3.2.2 → data_factory_utils-3.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "data-factory-utils"
-version = "3.2.2"
+version = "3.3.0"
 description = "Utility functions for interacting with data factories."
 readme = "README.md"
 requires-python = ">=3.12"

{data_factory_utils-3.2.2 → data_factory_utils-3.3.0}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "data-factory-utils"
-version = "3.2.2"
+version = "3.3.0"
 description = "Utility functions for interacting with data factories."
 readme = "README.md"
 requires-python = ">=3.12"

{data_factory_utils-3.2.2 → data_factory_utils-3.3.0}/src/data_factory_utils/athena.py

@@ -200,24 +200,23 @@ class AthenaQuery:
         # Get the raw query results from the API
         responses = self._get_responses()
 
-        # Parse the raw row data from the results
-        rows = []
+        column_names: list[str | None] | None = None
+        columns: list[list[str | None]] = []
+
         for response in responses:
             for row in response["ResultSet"]["Rows"]:
                 row_data = [column.get("VarCharValue") for column in row["Data"]]
-                rows.append(row_data)
-
-        if len(rows) == 0:
-            return pl.DataFrame({})
 
-        # The first row returned by the response is the names of the columns.
-        column_names, values_in_rows = rows[0], rows[1:]
+                # The first row returned by the response is the names of the columns.
+                if column_names is None:
+                    column_names = row_data
+                    columns = [[] for _ in column_names]
+                    continue
 
-        if len(values_in_rows) == 0:
-            return pl.DataFrame(data={col: [] for col in column_names})
+                for column, value in zip(columns, row_data, strict=True):
+                    column.append(value)
 
-        # Convert a list of rows into a list of columns
-        values_in_columns = list(zip(*values_in_rows, strict=True))
+        if column_names is None:
+            return pl.DataFrame({})
 
-        # Package it up into a dictionary of columns
-        return pl.DataFrame(dict(zip(column_names, values_in_columns, strict=True)))
+        return pl.DataFrame(dict(zip(column_names, columns, strict=True)))