chore: add compile_filter (#1817)

Author: chelsea-lin

bigframes/core/compile/sqlglot/compiler.py

@@ -211,6 +211,13 @@ def compile_projection(
         )
         return child.project(projected_cols)
 
+    @_compile_node.register
+    def compile_filter(
+        self, node: nodes.FilterNode, child: ir.SQLGlotIR
+    ) -> ir.SQLGlotIR:
+        condition = scalar_compiler.compile_scalar_expression(node.predicate)
+        return child.filter(condition)
+
     @_compile_node.register
     def compile_concat(
         self, node: nodes.ConcatNode, *children: ir.SQLGlotIR

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -99,3 +99,10 @@ def compile_addop(op: ops.AddOp, left: TypedExpr, right: TypedExpr) -> sge.Expre
 
     # Numerical addition
     return sge.Add(this=left.expr, expression=right.expr)
+
+
+def compile_ge(
+    op: ops.ge_op, left: TypedExpr, right: TypedExpr  # type: ignore[valid-type]
+) -> sge.Expression:
+
+    return sge.GTE(this=left.expr, expression=right.expr)

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -250,6 +250,16 @@ def project(
         new_expr = self._encapsulate_as_cte().select(*projected_cols_expr, append=True)
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
+    def filter(
+        self,
+        condition: sge.Expression,
+    ) -> SQLGlotIR:
+        """Filters the query with the given condition."""
+        new_expr = self._encapsulate_as_cte()
+        return SQLGlotIR(
+            expr=new_expr.where(condition, append=False), uid_gen=self.uid_gen
+        )
+
     def insert(
         self,
         destination: bigquery.TableReference,

tests/unit/core/compile/sqlglot/snapshots/test_compile_filter/test_compile_filter/out.sql

@@ -0,0 +1,25 @@
+WITH `bfcte_0` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_1` AS (
+  SELECT
+    *,
+    `bfcol_1` AS `bfcol_5`,
+    `bfcol_1` AS `bfcol_6`,
+    `bfcol_0` AS `bfcol_7`,
+    `bfcol_1` >= 1 AS `bfcol_8`
+  FROM `bfcte_0`
+), `bfcte_2` AS (
+  SELECT
+    *
+  FROM `bfcte_1`
+  WHERE
+    `bfcol_8`
+)
+SELECT
+  `bfcol_5` AS `rowindex`,
+  `bfcol_6` AS `rowindex_1`,
+  `bfcol_7` AS `int64_col`
+FROM `bfcte_2`

tests/unit/core/compile/sqlglot/test_compile_filter.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+import bigframes.pandas as bpd
+
+pytest.importorskip("pytest_snapshot")
+
+
+def test_compile_filter(scalars_types_df: bpd.DataFrame, snapshot):
+    bf_df = scalars_types_df[["rowindex", "int64_col"]]
+    bf_filter = bf_df[bf_df["rowindex"] >= 1]
+    snapshot.assert_match(bf_filter.sql, "out.sql")