fix: prune extra columns in Block.__init__ to prevent unbound variables

Author: shuoweil

bigframes/core/blocks.py

@@ -161,7 +161,7 @@ def __init__(
         self._expr = self._normalize_expression(expr, self._index_columns)
 
         # Calculate value_columns after normalizing expression
-        actual_value_columns = [
+        all_value_columns = [
             column
             for column in self._expr.column_ids
             if column not in self.index_columns
@@ -173,31 +173,20 @@ def __init__(
             if isinstance(column_labels, pd.Index)
             else pd.Index(column_labels)
         )
-
-        # 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(actual_value_columns)}) and "
-                f"'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
-
+        # Adjust column_labels and value_columns to match
+        if len(all_value_columns) > len(self._column_labels):
+            # More columns than labels: Drop the extra columns (assumed to be internal/garbage)
+            self._value_columns = all_value_columns[: len(self._column_labels)]
+            # Prune the expression to remove hidden columns
+            self._expr = self._expr.select_columns(
+                [*self.index_columns, *self._value_columns]
+            )
+        elif len(all_value_columns) < len(self._column_labels):
+            # Fewer columns than labels: Truncate labels
+            self._value_columns = all_value_columns
+            self._column_labels = self._column_labels[: len(self._value_columns)]
+        else:
+            self._value_columns = all_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]] = {