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
@@ -174,29 +174,28 @@ def __init__(
             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)]
+        # DEBUG
+        print(
+            f"DEBUG: Block init. all_value_columns={len(all_value_columns)}, column_labels={len(self._column_labels)}"
+        )
 
-        # 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"
+        # Adjust column_labels and value_columns to match
+        if len(all_value_columns) > len(self._column_labels):
+            print(
+                f"DEBUG: Pruning {len(all_value_columns) - len(self._column_labels)} extra columns"
             )
-
-        # Update value_columns property to use actual_value_columns
-        self._value_columns = actual_value_columns
+            # 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)

bigframes/core/expression.py

@@ -316,6 +316,10 @@ def bind_refs(
         if self.id in bindings.keys():
             return bindings[self.id]
         elif not allow_partial_bindings:
+            # DEBUG: Print details about the failure
+            print(f"DEBUG: bind_refs failed for {self.id}")
+            print(f"DEBUG: available bindings: {list(bindings.keys())}")
+            # print(f"DEBUG: full expression: {self}")
             raise ValueError(f"Variable {self.id} remains unbound")
         return self