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

tests/system/small/test_session.py

@@ -309,64 +309,6 @@ 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]
-            or 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_with_columns_filter(session, scalars_table_id: str):
-    cols = ["int64_too", "string_col", "date_col"]
-    df = session.read_gbq(scalars_table_id, columns=cols)
-    assert list(df.columns) == cols
-
-
 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)

tests/unit/session/test_session.py

@@ -57,3 +57,60 @@ def test_session_init_fails_with_no_project():
                 credentials=mock.Mock(spec=google.auth.credentials.Credentials)
             )
         )
+
+
+@pytest.mark.parametrize(
+    ("query_or_table", "columns", "filters", "expected_output"),
+    [
+        pytest.param(
+            """SELECT
+                rowindex,
+                string_col,
+            FROM `test_table` AS t
+            """,
+            [],
+            [("rowindex", "<", 4), ("string_col", "==", "Hello, World!")],
+            """SELECT * FROM (SELECT
+                rowindex,
+                string_col,
+            FROM `test_table` AS t
+            ) AS sub WHERE `rowindex` < 4 AND `string_col` = 'Hello, World!'""",
+            id="query_input",
+        ),
+        pytest.param(
+            "test_table",
+            [],
+            [("date_col", ">", "2022-10-20")],
+            "SELECT * FROM test_table AS sub WHERE `date_col` > '2022-10-20'",
+            id="table_input",
+        ),
+        pytest.param(
+            "test_table",
+            ["row_index", "string_col"],
+            [
+                (("rowindex", "not in", [0, 6]),),
+                (("string_col", "in", ["Hello, World!", "こんにちは"]),),
+            ],
+            (
+                "SELECT `row_index`, `string_col` FROM test_table AS sub WHERE "
+                "`rowindex` NOT IN (0, 6) OR `string_col` IN ('Hello, World!', "
+                "'こんにちは')"
+            ),
+            id="or_operation",
+        ),
+        pytest.param(
+            "test_table",
+            [],
+            ["date_col", ">", "2022-10-20"],
+            None,
+            marks=pytest.mark.xfail(
+                raises=ValueError,
+            ),
+            id="raise_error",
+        ),
+    ],
+)
+def test_read_gbq_with_filters(query_or_table, columns, filters, expected_output):
+    session = resources.create_bigquery_session()
+    query = session._filters_to_query(query_or_table, columns, filters)
+    assert query == expected_output

third_party/bigframes_vendored/pandas/io/gbq.py

@@ -7,20 +7,8 @@
 
 from bigframes import constants
 
-FiltersType = (
-    Iterable[
-        Union[
-            Tuple[str, Literal["in", "not in", "<", "<=", "==", "!=", ">=", ">"], Any],
-            Iterable[
-                Tuple[
-                    str,
-                    Literal["in", "not in", "<", "<=", "==", "!=", ">=", ">"],
-                    Any,
-                ]
-            ],
-        ]
-    ],
-)
+FilterType = Tuple[str, Literal["in", "not in", "<", "<=", "==", "!=", ">=", ">"], Any]
+FiltersType = Iterable[Union[FilterType, Iterable[FilterType]]]
 
 
 class GBQIOMixin:
@@ -100,15 +88,15 @@ def read_gbq(
             max_results (Optional[int], default None):
                 If set, limit the maximum number of rows to fetch from the
                 query results.
-            columns(Iterable[str], default ()): If not empty, only these columns
+            columns (Iterable[str], default ()): If not empty, only these columns
                 will be read from table.
-            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.
+            filters (Iterable[Iterable[[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.