Merge pull request #1692 from no23reason/dho/cq-2677-resources-2

perf(gooddata-pandas): speed up and slim down DataFrame conversion

Author: no23reason

packages/gooddata-pandas/src/gooddata_pandas/arrow_convertor.py

@@ -4,14 +4,16 @@
 import logging
 from typing import Callable
 
+import numpy
 import orjson
 import pandas
-from gooddata_sdk.type_converter import AttributeConverterStore
+from gooddata_sdk.type_converter import AttributeConverterStore, DateConverter, DatetimeConverter
 
 from gooddata_pandas.arrow_types import TypesMapper
 
 try:
     import pyarrow as pa
+    import pyarrow.compute as pc
 except ImportError as _exc:
     raise ImportError(
         "pyarrow is required for Arrow support. Install it with: pip install gooddata-pandas[arrow]"
@@ -96,9 +98,30 @@ def convert_label_values(label_id: str, values: list, model_labels: dict) -> lis
     Returns:
         Converted list, or the original *values* object when no conversion is needed.
     """
+    # pick the converter for this label's granularity (None for plain text attributes)
     converter = _get_date_converter_for_label(label_id, model_labels)
     if converter is None:
         return values
+
+    if isinstance(converter, (DateConverter, DatetimeConverter)):
+        # Date/datetime granularity is costly, so vectorize it in three steps:
+
+        # 1) parse each raw string into a datetime.date/datetime with the cheap
+        #    per-value to_type() (handles partial dates like "2023" or "2023-01");
+        #    keep None as None. This step stays per-value so that pandas
+        #    does not re-infer the date format from the raw strings in step 2.
+        typed = [converter.to_type(v) if v is not None else None for v in values]
+
+        # 2) convert the whole column to Timestamps in one vectorized call (this
+        #    single call replaces N per-value ones; None becomes NaT here).
+        converted = pandas.to_datetime(typed)
+
+        # 3) rebuild the list, restoring None wherever the input was None so that NaT
+        #    does not leak into the output.
+        return [None if orig is None else c for orig, c in zip(values, converted)]
+
+    # WEEK / QUARTER (StringConverter) and integer granularities are cheap per value
+    # and could change type subtly if batched, so convert them one by one (None kept).
     return [converter.to_external_type(v) if v is not None else None for v in values]
 
 
@@ -420,16 +443,23 @@ def reorder_grand_totals(
         return table
     if _COL_ROW_TYPE not in table.schema.names:
         return table
-    row_type_vals = table.column(_COL_ROW_TYPE).to_pylist()
-    grand_mask = pa.array([v == 2 for v in row_type_vals], type=pa.bool_())
-    grand_total_rows = table.filter(grand_mask)
+
+    # uses pyarrow.compute to run these in a vectorized way.
+    # the fill_null ensures that nulls are still marked as non-totals.
+    # (pyarrow.compute functions are generated at runtime, so ty cannot see them.)
+    row_type_col = table.column(_COL_ROW_TYPE)
+    is_grand_total = pc.fill_null(pc.equal(row_type_col, 2), False)  # ty: ignore[unresolved-attribute]
+    grand_total_rows = table.filter(is_grand_total)
     if grand_total_rows.num_rows == 0:
         return table
-    data_and_sub_rows = table.filter(pa.array([v != 2 for v in row_type_vals], type=pa.bool_()))
-    return pa.concat_tables([grand_total_rows, data_and_sub_rows])
+
+    not_grand_total = pc.fill_null(pc.not_equal(row_type_col, 2), True)  # ty: ignore[unresolved-attribute]
+    return pa.concat_tables([grand_total_rows, table.filter(not_grand_total)])
 
 
-def compute_column_totals_indexes(table: pa.Table, execution_dims: list) -> list[list[int]]:
+def compute_column_totals_indexes(
+    table: pa.Table, execution_dims: list, schema_meta: dict | None = None
+) -> list[list[int]]:
     """
     Compute column_totals_indexes compatible with DataFrameMetadata from an Arrow table.
 
@@ -445,7 +475,9 @@ def compute_column_totals_indexes(table: pa.Table, execution_dims: list) -> list
         grand_total_* fields become output rows and are already covered by
         compute_row_totals_indexes.  Returns [] in that case.
     """
-    schema_meta = _parse_schema_metadata(table)
+    if schema_meta is None:
+        schema_meta = _parse_schema_metadata(table)
+
     xtab_meta = schema_meta[_META_XTAB]
     is_transposed = schema_meta[_META_VIEW]["isTransposed"]
 
@@ -518,7 +550,9 @@ def _label_ids_in_dim(dim: dict) -> set:
     return result
 
 
-def compute_row_totals_indexes(table: pa.Table, execution_dims: list) -> list[list[int]]:
+def compute_row_totals_indexes(
+    table: pa.Table, execution_dims: list, schema_meta: dict | None = None
+) -> list[list[int]]:
     """
     Compute row_totals_indexes compatible with DataFrameMetadata from an Arrow table.
 
@@ -534,7 +568,9 @@ def compute_row_totals_indexes(table: pa.Table, execution_dims: list) -> list[li
         Total rows are grand_total_N fields.  A field is total at level j only when
         j >= len(gdc["label_values"]), i.e. that label level is being aggregated.
     """
-    schema_meta = _parse_schema_metadata(table)
+    if schema_meta is None:
+        schema_meta = _parse_schema_metadata(table)
+
     xtab_meta = schema_meta[_META_XTAB]
     is_transposed = schema_meta[_META_VIEW]["isTransposed"]
 
@@ -608,8 +644,8 @@ def _label_ids_in_dim(dim: dict) -> set:
     else:
         # Output rows are Arrow rows; every total row (row_type != 0) is listed
         # in the total-indexes for every attribute level.
-        row_types = _get_row_types(table)
-        total_row_idxs = [i for i, rt in enumerate(row_types) if rt != 0]
+        row_types = table.column(_COL_ROW_TYPE).to_numpy(zero_copy_only=False)
+        total_row_idxs = numpy.nonzero(row_types != 0)[0].tolist()
 
         result = []
         for header in row_dim.get("headers", []):
@@ -643,30 +679,47 @@ def _compute_primary_labels_from_inline(
     """
     result: dict[int, dict[str, str]] = {}
     label_meta = xtab_meta.get("labelMetadata", {})
-    row_types = _get_row_types(table)
-    data_row_mask = [rt == 0 for rt in row_types]
 
-    for j, ref in enumerate(label_refs):
+    # Project to only the columns this function reads - __row_type plus the label
+    # and primary-label columns - before filtering. Filtering the whole table would
+    # also copy every metric column for the data rows even though only these
+    # attribute columns are ever read below.
+    # table.select is zero-copy, so the subsequent filter copies just these columns.
+    needed_cols = [_COL_ROW_TYPE]
+    for ref in label_refs:
         info = label_meta.get(ref, {})
         label_id = label_ref_to_id.get(ref, info.get("labelId", ""))
         primary_label_id = info.get("primaryLabelId", label_id)
+        for col in (label_id, primary_label_id):
+            if col in table.schema.names and col not in needed_cols:
+                needed_cols.append(col)
+
+    projected = table.select(needed_cols)
+
+    # Extract the data rows (row_type == 0) once and reuse. Only filter when totals
+    # are actually present - otherwise the projected table already is the data rows.
+    # (pyarrow.compute functions are generated at runtime, so ty cannot see them.)
+    row_type_col = projected.column(_COL_ROW_TYPE)
+    has_total_rows = pc.any(pc.not_equal(row_type_col, 0)).as_py()  # ty: ignore[unresolved-attribute]
+    data_row_mask = pc.equal(row_type_col, 0)  # ty: ignore[unresolved-attribute]
+    data_rows = projected.filter(data_row_mask) if has_total_rows else projected
 
-        display_vals = table.column(label_id).to_pylist()
+    for j, ref in enumerate(label_refs):
+        info = label_meta.get(ref, {})
+        label_id = label_ref_to_id.get(ref, info.get("labelId", ""))
+        primary_label_id = info.get("primaryLabelId", label_id)
 
-        if label_id == primary_label_id:
+        if label_id == primary_label_id or primary_label_id not in table.schema.names:
+            # identity (or fallback when the primary column is absent): map each
+            # distinct display value to itself.
             mapping: dict[str, str] = {
-                v: v for v, is_data in zip(display_vals, data_row_mask) if is_data and isinstance(v, str)
-            }
-        elif primary_label_id in table.schema.names:
-            primary_vals = table.column(primary_label_id).to_pylist()
-            mapping = {
-                p: d
-                for p, d, is_data in zip(primary_vals, display_vals, data_row_mask)
-                if is_data and isinstance(p, str) and isinstance(d, str)
+                v: v for v in data_rows.column(label_id).unique().to_pylist() if isinstance(v, str)
             }
         else:
-            # Fallback: identity (primary label data not present in table)
-            mapping = {v: v for v, is_data in zip(display_vals, data_row_mask) if is_data and isinstance(v, str)}
+            # primary != display: map each primary value to its display value
+            primary_vals = data_rows.column(primary_label_id).to_pylist()
+            display_vals = data_rows.column(label_id).to_pylist()
+            mapping = {p: d for p, d in zip(primary_vals, display_vals) if isinstance(p, str) and isinstance(d, str)}
 
         result[j] = mapping
     return result
@@ -709,6 +762,7 @@ def _compute_primary_labels_from_fields(
 
 def compute_primary_labels(
     table: pa.Table,
+    schema_meta: dict | None = None,
 ) -> tuple[dict[int, dict[str, str]], dict[int, dict[str, str]]]:
     """
     Compute primary_labels_from_index and primary_labels_from_columns from an Arrow table.
@@ -719,7 +773,9 @@ def compute_primary_labels(
     Returns:
         (primary_labels_from_index, primary_labels_from_columns)
     """
-    schema_meta = _parse_schema_metadata(table)
+    if schema_meta is None:
+        schema_meta = _parse_schema_metadata(table)
+
     xtab_meta = schema_meta[_META_XTAB]
     is_transposed = schema_meta[_META_VIEW]["isTransposed"]
 
@@ -750,6 +806,7 @@ def convert_arrow_table_to_dataframe(
     types_mapper: TypesMapper = TypesMapper.DEFAULT,
     custom_mapping: dict | None = None,
     label_overrides: dict | None = None,
+    schema_meta: dict | None = None,
 ) -> pandas.DataFrame:
     """
     Convert a pyarrow Table returned by the GoodData /binary execution endpoint
@@ -800,7 +857,9 @@ def convert_arrow_table_to_dataframe(
     else:
         raise ValueError("Unknown types_mapper value")
 
-    schema_meta = _parse_schema_metadata(table)
+    if schema_meta is None:
+        schema_meta = _parse_schema_metadata(table)
+
     xtab_meta = schema_meta[_META_XTAB]
     model_meta = schema_meta[_META_MODEL]
     is_transposed = schema_meta[_META_VIEW]["isTransposed"]

packages/gooddata-pandas/src/gooddata_pandas/data_access.py

@@ -25,7 +25,7 @@
     _str_to_obj_id,
     _to_attribute,
     _to_item,
-    _typed_attribute_value,
+    _typed_attribute_values,
     get_catalog_attributes_for_extract,
 )
 
@@ -340,22 +340,21 @@ def _find_attribute(attributes: list[CatalogAttribute], id_obj: IdObjType) -> Un
     return None
 
 
-def _typed_result(attributes: list[CatalogAttribute], attribute: Attribute, result_values: list[Any]) -> list[Any]:
+def _resolve_catalog_attribute(attributes: list[CatalogAttribute], attribute: Attribute) -> CatalogAttribute:
     """
-    Internal function to convert result_values to proper data types.
+    Find the CatalogAttribute matching the given execution attribute.
 
     Args:
         attributes (list[CatalogAttribute]): The catalog of attributes.
-        attribute (Attribute): The attribute for which the typed result will be computed.
-        result_values (list[Any]): A list of raw values.
+        attribute (Attribute): The execution attribute to resolve.
 
     Returns:
-        list[Any]: A list of converted values with proper data types.
+        CatalogAttribute: The matching catalog attribute.
     """
     catalog_attribute = _find_attribute(attributes, attribute.label)
     if catalog_attribute is None:
         raise ValueError(f"Unable to find attribute {attribute.label} in catalog")
-    return [_typed_attribute_value(catalog_attribute, value) for value in result_values]
+    return catalog_attribute
 
 
 def _extract_from_attributes_and_maybe_metrics(
@@ -399,20 +398,28 @@ def _extract_from_attributes_and_maybe_metrics(
     index_to_attribute = {index_name: exec_def.attributes[i] for index_name, i in safe_index_to_attr_idx.items()}
     col_to_attribute = {col: exec_def.attributes[i] for col, i in col_to_attr_idx.items()}
 
+    # resolve the matching CatalogAttribute for each index / attribute column once:
+    # it does not change during the batch iteration
+    index_to_catalog_attribute = {
+        index_name: _resolve_catalog_attribute(attributes, attribute)
+        for index_name, attribute in index_to_attribute.items()
+    }
+    col_to_catalog_attribute = {
+        col: _resolve_catalog_attribute(attributes, attribute) for col, attribute in col_to_attribute.items()
+    }
+
     # datastructures to return
     index: dict[str, list[Any]] = {idx_name: [] for idx_name in safe_index_to_attr_idx}
     data: dict[str, list[Any]] = {col: [] for col in cols}
 
     while True:
         for idx_name in index:
             rs = result.get_all_header_values(attribute_dim, safe_index_to_attr_idx[idx_name])
-            attribute = index_to_attribute[idx_name]
-            index[idx_name] += _typed_result(attributes, attribute, rs)
+            index[idx_name] += _typed_attribute_values(index_to_catalog_attribute[idx_name], rs)
         for col in cols:
             if col in col_to_attr_idx:
                 rs = result.get_all_header_values(attribute_dim, col_to_attr_idx[col])
-                attribute = col_to_attribute[col]
-                data[col] += _typed_result(attributes, attribute, rs)
+                data[col] += _typed_attribute_values(col_to_catalog_attribute[col], rs)
             elif col_to_metric_idx[col] < len(result.data):
                 data[col] += result.data[col_to_metric_idx[col]]
         if result.is_complete(attribute_dim):
@@ -455,11 +462,11 @@ def _extract_from_arrow(
     metric_dim_idx_to_field = build_metric_field_index(table)
     model_labels = read_model_labels(table)
 
-    data: dict[str, list] = {}
+    data: dict[str, Any] = {}
     for col in cols:
         if col in col_to_metric_idx:
             field_name = metric_dim_idx_to_field[col_to_metric_idx[col]]
-            data[col] = table.column(field_name).to_pylist()
+            data[col] = table.column(field_name).to_numpy(zero_copy_only=False)
         else:
             attr = exec_def.attributes[col_to_attr_idx[col]]
             label_id = attr.label.id

packages/gooddata-pandas/src/gooddata_pandas/dataframe.py

@@ -25,6 +25,7 @@
 
 try:
     from gooddata_pandas.arrow_convertor import (
+        _parse_schema_metadata,
         compute_column_totals_indexes,
         compute_primary_labels,
         compute_row_totals_indexes,
@@ -469,16 +470,18 @@ def _table_to_df_and_metadata(
         call site.
         """
         table = reorder_grand_totals(table, grand_totals_position)
+        schema_meta = _parse_schema_metadata(table)
         df = convert_arrow_table_to_dataframe(
             table,
             self_destruct=self._arrow_config.self_destruct,
             types_mapper=self._arrow_config.types_mapper,
             custom_mapping=self._arrow_config.custom_mapping,
             label_overrides=label_overrides,
+            schema_meta=schema_meta,
         )
-        row_totals_indexes = compute_row_totals_indexes(table, exec_response.dimensions)
-        column_totals_indexes = compute_column_totals_indexes(table, exec_response.dimensions)
-        primary_labels_from_index, primary_labels_from_columns = compute_primary_labels(table)
+        row_totals_indexes = compute_row_totals_indexes(table, exec_response.dimensions, schema_meta=schema_meta)
+        column_totals_indexes = compute_column_totals_indexes(table, exec_response.dimensions, schema_meta=schema_meta)
+        primary_labels_from_index, primary_labels_from_columns = compute_primary_labels(table, schema_meta=schema_meta)
         metadata = DataFrameMetadata(
             row_totals_indexes=row_totals_indexes,
             column_totals_indexes=column_totals_indexes,
@@ -588,22 +591,28 @@ def for_arrow_table(
             label_overrides = {}
 
         table = reorder_grand_totals(table, grand_totals_position)
+        # parse the schema metadata once and share it across the sibling functions
+        # below; each of them would otherwise re-parse it from the table
+        schema_meta = _parse_schema_metadata(table)
         df = convert_arrow_table_to_dataframe(
             table,
             self_destruct=self._arrow_config.self_destruct,
             types_mapper=self._arrow_config.types_mapper,
             custom_mapping=self._arrow_config.custom_mapping,
             label_overrides=label_overrides,
+            schema_meta=schema_meta,
         )
         row_totals_indexes = (
-            compute_row_totals_indexes(table, execution_response.dimensions) if execution_response is not None else []
+            compute_row_totals_indexes(table, execution_response.dimensions, schema_meta=schema_meta)
+            if execution_response is not None
+            else []
         )
         column_totals_indexes = (
-            compute_column_totals_indexes(table, execution_response.dimensions)
+            compute_column_totals_indexes(table, execution_response.dimensions, schema_meta=schema_meta)
             if execution_response is not None
             else []
         )
-        primary_labels_from_index, primary_labels_from_columns = compute_primary_labels(table)
+        primary_labels_from_index, primary_labels_from_columns = compute_primary_labels(table, schema_meta=schema_meta)
         metadata = DataFrameMetadata(
             row_totals_indexes=row_totals_indexes,
             column_totals_indexes=column_totals_indexes,