Skip to content

Commit edd34f0

Browse files
fix agg case, flexible col names
1 parent 4ff8a57 commit edd34f0

File tree

3 files changed

+32
-1
lines changed

3 files changed

+32
-1
lines changed

bigframes/core/nodes.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1568,6 +1568,10 @@ class ExplodeNode(UnaryNode):
15681568
# Offsets are generated only if this is non-null
15691569
offsets_col: Optional[identifiers.ColumnId] = None
15701570

1571+
def _validate(self):
1572+
for col in self.column_ids:
1573+
assert col.id in self.child.ids
1574+
15711575
@property
15721576
def row_preserving(self) -> bool:
15731577
return False
@@ -1646,6 +1650,10 @@ class ResultNode(UnaryNode):
16461650
limit: Optional[int] = None
16471651
# TODO: CTE definitions
16481652

1653+
def _validate(self):
1654+
for ref, name in self.output_cols:
1655+
assert ref.id in self.child.ids
1656+
16491657
@property
16501658
def node_defined_ids(self) -> Tuple[identifiers.ColumnId, ...]:
16511659
return ()

bigframes/core/rewrite/select_pullup.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,27 @@ def pull_up_select_unary(node: nodes.UnaryNode) -> nodes.BigFrameNode:
6565
nodes.SelectionNode, child.replace_child(pushed_down_node)
6666
)
6767
return pulled_up_select
68-
elif isinstance(node, (nodes.SelectionNode, nodes.ResultNode, nodes.AggregateNode)):
68+
elif isinstance(
69+
node,
70+
(
71+
nodes.SelectionNode,
72+
nodes.ResultNode,
73+
),
74+
):
6975
return node.remap_refs(
7076
{id: ref.id for ref, id in child.input_output_pairs}
7177
).replace_child(child.child)
78+
elif isinstance(node, nodes.AggregateNode):
79+
pushed_down_agg = node.remap_refs(
80+
{id: ref.id for ref, id in child.input_output_pairs}
81+
).replace_child(child.child)
82+
new_selection = tuple(
83+
nodes.AliasedRef.identity(id).remap_refs(
84+
{id: ref.id for ref, id in child.input_output_pairs}
85+
)
86+
for id in node.ids
87+
)
88+
return nodes.SelectionNode(pushed_down_agg, new_selection)
7289
elif isinstance(node, nodes.ExplodeNode):
7390
pushed_down_node = node.remap_refs(
7491
{id: ref.id for ref, id in child.input_output_pairs}

bigframes/core/utils.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,12 @@ def label_to_identifier(label: typing.Hashable, strict: bool = False) -> str:
148148
# first character must be letter or underscore
149149
identifier = "_" + identifier
150150

151+
else:
152+
# even with flexible column names, there are constraints
153+
# Convert illegal characters
154+
# See: https://cloud.google.com/bigquery/docs/schemas#flexible-column-names
155+
identifier = re.sub(r"[!\"$\(\)\*\,\./;\?@[\]^`{}~]", "_", identifier)
156+
151157
# Except in special circumstances (true anonymous query results tables),
152158
# field names are not allowed to start with these (case-insensitive)
153159
# prefixes.

0 commit comments

Comments
 (0)