fix: handle value_columns and column_labels mismatch in Block

Author: shuoweil

bigframes/core/blocks.py

@@ -159,20 +159,48 @@ def __init__(
             else tuple([None for _ in index_columns])
         )
         self._expr = self._normalize_expression(expr, self._index_columns)
+
+        # FIX: Calculate value_columns after normalizing expression
+        actual_value_columns = [
+            column
+            for column in self._expr.column_ids
+            if column not in self.index_columns
+        ]
+
         # Use pandas index to more easily replicate column indexing, especially for hierarchical column index
         self._column_labels = (
             column_labels.copy()
             if isinstance(column_labels, pd.Index)
             else pd.Index(column_labels)
         )
-        if len(self.value_columns) != len(self._column_labels):
+
+        # FIX: Adjust column_labels if needed to match actual_value_columns
+        if len(actual_value_columns) != len(self._column_labels):
+            # If we have more value columns than labels, extend with None
+            if len(actual_value_columns) > len(self._column_labels):
+                additional_labels = [None] * (
+                    len(actual_value_columns) - len(self._column_labels)
+                )
+                self._column_labels = self._column_labels.append(
+                    pd.Index(additional_labels)
+                )
+            # If we have fewer value columns than labels, truncate
+            else:
+                self._column_labels = self._column_labels[: len(actual_value_columns)]
+
+        # Re-validate after adjustment
+        if len(actual_value_columns) != len(self._column_labels):
             raise ValueError(
-                f"'value_columns' (size {len(self.value_columns)}) and 'column_labels' (size {len(self._column_labels)}) must have equal length"
+                f"'value_columns' (size {len(actual_value_columns)}) and 'column_labels' (size {len(self._column_labels)}) must have equal length"
             )
+
+        # Update value_columns property to use actual_value_columns
+        self._value_columns = actual_value_columns
+
         # col_id -> [stat_name -> scalar]
         # TODO: Preserve cache under safe transforms (eg. drop column, reorder)
         self._stats_cache: dict[str, dict[str, typing.Any]] = {
-            col_id: {} for col_id in self.value_columns
+            col_id: {} for col_id in self._value_columns
         }
         # TODO(kemppeterson) Add a cache for corr to parallel the single-column stats.
 
@@ -285,11 +313,15 @@ def index_columns(self) -> Sequence[str]:
     @property
     def value_columns(self) -> Sequence[str]:
         """All value columns, mutually exclusive with index columns."""
-        return [
-            column
-            for column in self._expr.column_ids
-            if column not in self.index_columns
-        ]
+        return getattr(
+            self,
+            "_value_columns",
+            [
+                column
+                for column in self._expr.column_ids
+                if column not in self.index_columns
+            ],
+        )
 
     @property
     def column_labels(self) -> pd.Index: