feat: Add filters and columns arguments to read_gbq for enhanced data querying

bigframes/pandas/__init__.py

@@ -486,6 +486,7 @@ def read_gbq(
     index_col: Iterable[str] | str = (),
     col_order: Iterable[str] = (),
     max_results: Optional[int] = None,
+    filters: Optional[List[Tuple]] = None,
 ) -> bigframes.dataframe.DataFrame:
     _set_default_session_location_if_possible(query_or_table)
     return global_session.with_default_session(
@@ -494,6 +495,7 @@ def read_gbq(
         index_col=index_col,
         col_order=col_order,
         max_results=max_results,
+        filters=filters,
     )
 
 

bigframes/session/__init__.py

@@ -284,9 +284,11 @@ def read_gbq(
         index_col: Iterable[str] | str = (),
         col_order: Iterable[str] = (),
         max_results: Optional[int] = None,
+        filters: Optional[List[Tuple]] = None
         # Add a verify index argument that fails if the index is not unique.
     ) -> dataframe.DataFrame:
         # TODO(b/281571214): Generate prompt to show the progress of read_gbq.
+        query_or_table = self._filters_to_query(query_or_table, filters)
         if _is_query(query_or_table):
             return self._read_gbq_query(
                 query_or_table,
@@ -307,6 +309,72 @@ def read_gbq(
                 api_name="read_gbq",
             )
 
+    def _filters_to_query(self, query_or_table, filters):
+        """Convert filters to query"""
+
+        if (filters is None) or (len(filters) == 0):
+            return query_or_table
+
+        valid_operators = ["IN", "NOT IN", "=", ">", "<", ">=", "<=", "!="]
+
+        sub_query = (
+            f"({query_or_table})" if _is_query(query_or_table) else query_or_table
+        )
+
+        where_clause = ""
+        if filters:
+            if not isinstance(filters, list):
+                raise ValueError("Filters should be a list.")
+
+            if not (
+                all(isinstance(item, list) for item in filters)
+                or all(isinstance(item, tuple) for item in filters)
+            ):
+                raise ValueError(
+                    "All items in filters should be either all lists or all tuples."
+                )
+
+            if all(isinstance(sub_filter, tuple) for sub_filter in filters):
+                filters = [filters]
+
+            grouped_expressions = []
+            for group in filters:
+                if not isinstance(group, list):
+                    raise ValueError("Each filter group should be a list.")
+
+                group_expressions = []
+                for filter_item in group:
+                    if not isinstance(filter_item, tuple):
+                        raise ValueError("Each filter condition should be a tuple.")
+
+                    column, operator, value = filter_item
+                    operator = operator.upper()
+
+                    if operator not in valid_operators:
+                        raise ValueError(f"Operator {operator} is not valid.")
+
+                    if operator in ["IN", "NOT IN"]:
+                        if not isinstance(value, list):
+                            raise ValueError(
+                                f"Value for operator {operator} should be a list."
+                            )
+                        value_list = ", ".join(
+                            [f'"{v}"' if isinstance(v, str) else str(v) for v in value]
+                        )
+                        expression = f"{column} {operator} ({value_list})"
+                    else:
+                        value = f'"{value}"' if isinstance(value, str) else value
+                        expression = f"{column} {operator} {value}"
+                    group_expressions.append(expression)
+
+                grouped_expressions.append(" AND ".join(group_expressions))
+
+            where_clause = " WHERE " + " OR ".join(grouped_expressions)
+
+        full_query = f"SELECT * FROM {sub_query} AS sub{where_clause}"
+
+        return full_query
+
     def _query_to_destination(
         self,
         query: str,

tests/system/small/test_session.py

@@ -309,6 +309,58 @@ def test_read_gbq_w_script_no_select(session, dataset_id: str):
     assert df["statement_type"][0] == "SCRIPT"
 
 
+@pytest.mark.parametrize(
+    ("query_or_table", "filters", "validator"),
+    [
+        pytest.param(
+            """SELECT
+                rowindex,
+                string_col,
+            FROM `{scalars_table_id}` AS t
+            """,
+            [("rowindex", "<", 4), ("string_col", "=", "Hello, World!")],
+            lambda row: row["rowindex"] < 4 and row["string_col"] == "Hello, World!",
+            id="query_input",
+        ),
+        pytest.param(
+            "{scalars_table_id}",
+            [("date_col", ">", "2022-10-20")],
+            lambda row: pd.to_datetime(row["date_col"]) > pd.to_datetime("2022-10-20"),
+            id="table_input",
+        ),
+        pytest.param(
+            "{scalars_table_id}",
+            [
+                ("rowindex", "not in", [0, 6]),
+                ("string_col", "in", ["Hello, World!", "こんにちは"]),
+            ],
+            lambda row: row["rowindex"] not in [0, 6]
+            and row["string_col"] in ["Hello, World!", "こんにちは"],
+            id="or_operation",
+        ),
+        pytest.param(
+            "{scalars_table_id}",
+            ["date_col", ">", "2022-10-20"],
+            None,
+            marks=pytest.mark.xfail(
+                raises=ValueError,
+            ),
+            id="raise_error",
+        ),
+    ],
+)
+def test_read_gbq_with_filters(
+    session, scalars_table_id: str, query_or_table, filters, validator
+):
+    df = session.read_gbq(
+        query_or_table.format(scalars_table_id=scalars_table_id),
+        filters=filters,
+    )
+
+    for _, row in df.iterrows():
+        assert validator(row)
+
+
 def test_read_gbq_model(session, penguins_linear_model_name):
     model = session.read_gbq_model(penguins_linear_model_name)
     assert isinstance(model, bigframes.ml.linear_model.LinearRegression)

third_party/bigframes_vendored/pandas/io/gbq.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from typing import Iterable, Optional
+from typing import Iterable, List, Optional, Tuple
 
 from bigframes import constants
 
@@ -16,6 +16,7 @@ def read_gbq(
         index_col: Iterable[str] | str = (),
         col_order: Iterable[str] = (),
         max_results: Optional[int] = None,
+        filters: Optional[List[Tuple]] = None,
     ):
         """Loads a DataFrame from BigQuery.
 
@@ -83,6 +84,13 @@ def read_gbq(
             max_results (Optional[int], default None):
                 If set, limit the maximum number of rows to fetch from the
                 query results.
+            filters (List[Tuple], default []): To filter out data. Filter syntax:
+            [[(column, op, val), …],…] where op is [=, >, >=, <, <=, !=, in,
+            not in] The innermost tuples are transposed into a set of filters
+            applied through an AND operation. The outer list combines these sets
+            of filters through an OR operation. A single list of tuples can also
+            be used, meaning that no OR operation between set of filters is to be
+            conducted.
 
         Returns:
             bigframes.dataframe.DataFrame: A DataFrame representing results of the query or table.