Merge branch 'main' into main_chelsealin_bbqjson2

Author: chelsea-lin

bigframes/core/compile/sqlglot/compiler.py

@@ -244,6 +244,28 @@ def compile_join(
             joins_nulls=node.joins_nulls,
         )
 
+    @_compile_node.register
+    def compile_isin_join(
+        self, node: nodes.InNode, left: ir.SQLGlotIR, right: ir.SQLGlotIR
+    ) -> ir.SQLGlotIR:
+        conditions = (
+            typed_expr.TypedExpr(
+                scalar_compiler.compile_scalar_expression(node.left_col),
+                node.left_col.output_type,
+            ),
+            typed_expr.TypedExpr(
+                scalar_compiler.compile_scalar_expression(node.right_col),
+                node.right_col.output_type,
+            ),
+        )
+
+        return left.isin_join(
+            right,
+            indicator_col=node.indicator_col.sql,
+            conditions=conditions,
+            joins_nulls=node.joins_nulls,
+        )
+
     @_compile_node.register
     def compile_concat(
         self, node: nodes.ConcatNode, *children: ir.SQLGlotIR

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -336,6 +336,68 @@ def join(
 
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
+    def isin_join(
+        self,
+        right: SQLGlotIR,
+        indicator_col: str,
+        conditions: tuple[typed_expr.TypedExpr, typed_expr.TypedExpr],
+        joins_nulls: bool = True,
+    ) -> SQLGlotIR:
+        """Joins the current query with another SQLGlotIR instance."""
+        left_cte_name = sge.to_identifier(
+            next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
+        )
+
+        left_select = _select_to_cte(self.expr, left_cte_name)
+        # Prefer subquery over CTE for the IN clause's right side to improve SQL readability.
+        right_select = right.expr
+
+        left_ctes = left_select.args.pop("with", [])
+        right_ctes = right_select.args.pop("with", [])
+        merged_ctes = [*left_ctes, *right_ctes]
+
+        left_condition = typed_expr.TypedExpr(
+            sge.Column(this=conditions[0].expr, table=left_cte_name),
+            conditions[0].dtype,
+        )
+
+        new_column: sge.Expression
+        if joins_nulls:
+            right_table_name = sge.to_identifier(
+                next(self.uid_gen.get_uid_stream("bft_")), quoted=self.quoted
+            )
+            right_condition = typed_expr.TypedExpr(
+                sge.Column(this=conditions[1].expr, table=right_table_name),
+                conditions[1].dtype,
+            )
+            new_column = sge.Exists(
+                this=sge.Select()
+                .select(sge.convert(1))
+                .from_(sge.Alias(this=right_select.subquery(), alias=right_table_name))
+                .where(
+                    _join_condition(left_condition, right_condition, joins_nulls=True)
+                )
+            )
+        else:
+            new_column = sge.In(
+                this=left_condition.expr,
+                expressions=[right_select.subquery()],
+            )
+
+        new_column = sge.Alias(
+            this=new_column,
+            alias=sge.to_identifier(indicator_col, quoted=self.quoted),
+        )
+
+        new_expr = (
+            sge.Select()
+            .select(sge.Column(this=sge.Star(), table=left_cte_name), new_column)
+            .from_(sge.Table(this=left_cte_name))
+        )
+        new_expr.set("with", sge.With(expressions=merged_ctes))
+
+        return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
+
     def explode(
         self,
         column_names: tuple[str, ...],

bigframes/functions/_function_client.py

@@ -25,9 +25,11 @@
 import textwrap
 import types
 from typing import Any, cast, Optional, Sequence, Tuple, TYPE_CHECKING
+import warnings
 
 import requests
 
+import bigframes.exceptions as bfe
 import bigframes.formatting_helpers as bf_formatting
 import bigframes.functions.function_template as bff_template
 
@@ -482,10 +484,16 @@ def create_cloud_function(
                 function.service_config.max_instance_count = max_instance_count
             if vpc_connector is not None:
                 function.service_config.vpc_connector = vpc_connector
+                if vpc_connector_egress_settings is None:
+                    msg = bfe.format_message(
+                        "The 'vpc_connector_egress_settings' was not specified. Defaulting to 'private-ranges-only'.",
+                    )
+                    warnings.warn(msg, category=UserWarning)
+                    vpc_connector_egress_settings = "private-ranges-only"
                 if vpc_connector_egress_settings not in _VPC_EGRESS_SETTINGS_MAP:
                     raise bf_formatting.create_exception_with_feedback_link(
                         ValueError,
-                        f"'{vpc_connector_egress_settings}' not one of the supported vpc egress settings values: {list(_VPC_EGRESS_SETTINGS_MAP)}",
+                        f"'{vpc_connector_egress_settings}' is not one of the supported vpc egress settings values: {list(_VPC_EGRESS_SETTINGS_MAP)}",
                     )
                 function.service_config.vpc_connector_egress_settings = cast(
                     functions_v2.ServiceConfig.VpcConnectorEgressSettings,

bigframes/functions/_function_session.py

@@ -245,9 +245,9 @@ def remote_function(
         cloud_function_timeout: Optional[int] = 600,
         cloud_function_max_instances: Optional[int] = None,
         cloud_function_vpc_connector: Optional[str] = None,
-        cloud_function_vpc_connector_egress_settings: Literal[
-            "all", "private-ranges-only", "unspecified"
-        ] = "private-ranges-only",
+        cloud_function_vpc_connector_egress_settings: Optional[
+            Literal["all", "private-ranges-only", "unspecified"]
+        ] = None,
         cloud_function_memory_mib: Optional[int] = 1024,
         cloud_function_ingress_settings: Literal[
             "all", "internal-only", "internal-and-gclb"
@@ -514,6 +514,16 @@ def remote_function(
                 " For more details see https://cloud.google.com/functions/docs/securing/cmek#before_you_begin.",
             )
 
+        # A VPC connector is required to specify VPC egress settings.
+        if (
+            cloud_function_vpc_connector_egress_settings is not None
+            and cloud_function_vpc_connector is None
+        ):
+            raise bf_formatting.create_exception_with_feedback_link(
+                ValueError,
+                "cloud_function_vpc_connector must be specified before cloud_function_vpc_connector_egress_settings.",
+            )
+
         if cloud_function_ingress_settings is None:
             cloud_function_ingress_settings = "internal-only"
             msg = bfe.format_message(

bigframes/pandas/__init__.py

@@ -87,9 +87,9 @@ def remote_function(
     cloud_function_timeout: Optional[int] = 600,
     cloud_function_max_instances: Optional[int] = None,
     cloud_function_vpc_connector: Optional[str] = None,
-    cloud_function_vpc_connector_egress_settings: Literal[
-        "all", "private-ranges-only", "unspecified"
-    ] = "private-ranges-only",
+    cloud_function_vpc_connector_egress_settings: Optional[
+        Literal["all", "private-ranges-only", "unspecified"]
+    ] = None,
     cloud_function_memory_mib: Optional[int] = 1024,
     cloud_function_ingress_settings: Literal[
         "all", "internal-only", "internal-and-gclb"

bigframes/session/__init__.py

@@ -1509,9 +1509,9 @@ def remote_function(
         cloud_function_timeout: Optional[int] = 600,
         cloud_function_max_instances: Optional[int] = None,
         cloud_function_vpc_connector: Optional[str] = None,
-        cloud_function_vpc_connector_egress_settings: Literal[
-            "all", "private-ranges-only", "unspecified"
-        ] = "private-ranges-only",
+        cloud_function_vpc_connector_egress_settings: Optional[
+            Literal["all", "private-ranges-only", "unspecified"]
+        ] = None,
         cloud_function_memory_mib: Optional[int] = 1024,
         cloud_function_ingress_settings: Literal[
             "all", "internal-only", "internal-and-gclb"

noxfile.py

@@ -665,7 +665,7 @@ def prerelease(session: nox.sessions.Session, tests_path, extra_pytest_options=(
     session.install(
         "--upgrade",
         "-e",
-        "git+https://github.com/googleapis/python-bigquery-storage.git#egg=google-cloud-bigquery-storage",
+        "git+https://github.com/googleapis/google-cloud-python.git#egg=google-cloud-bigquery-storage&subdirectory=packages/google-cloud-bigquery-storage",
     )
     already_installed.add("google-cloud-bigquery-storage")
     session.install(

tests/system/large/functions/test_remote_function.py

@@ -1512,6 +1512,46 @@ def square_num(x):
         )
 
 
+@pytest.mark.flaky(retries=2, delay=120)
+def test_remote_function_no_vpc_connector(session):
+    def foo(x):
+        return x
+
+    with pytest.raises(
+        ValueError,
+        match="^cloud_function_vpc_connector must be specified before cloud_function_vpc_connector_egress_settings",
+    ):
+        session.remote_function(
+            input_types=[int],
+            output_type=int,
+            reuse=False,
+            cloud_function_service_account="default",
+            cloud_function_vpc_connector=None,
+            cloud_function_vpc_connector_egress_settings="all",
+            cloud_function_ingress_settings="all",
+        )(foo)
+
+
+@pytest.mark.flaky(retries=2, delay=120)
+def test_remote_function_wrong_vpc_egress_value(session):
+    def foo(x):
+        return x
+
+    with pytest.raises(
+        ValueError,
+        match="^'wrong-egress-value' is not one of the supported vpc egress settings values:",
+    ):
+        session.remote_function(
+            input_types=[int],
+            output_type=int,
+            reuse=False,
+            cloud_function_service_account="default",
+            cloud_function_vpc_connector="dummy-value",
+            cloud_function_vpc_connector_egress_settings="wrong-egress-value",
+            cloud_function_ingress_settings="all",
+        )(foo)
+
+
 @pytest.mark.parametrize(
     ("max_batching_rows"),
     [

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

@@ -85,7 +85,7 @@ def scalar_types_table_schema() -> typing.Sequence[bigquery.SchemaField]:
         bigquery.SchemaField("numeric_col", "NUMERIC"),
         bigquery.SchemaField("float64_col", "FLOAT"),
         bigquery.SchemaField("rowindex", "INTEGER"),
-        bigquery.SchemaField("rowindex_2", "INTEGER"),
+        bigquery.SchemaField("rowindex_2", "INTEGER", mode="REQUIRED"),
         bigquery.SchemaField("string_col", "STRING"),
         bigquery.SchemaField("time_col", "TIME"),
         bigquery.SchemaField("timestamp_col", "TIMESTAMP"),

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

@@ -0,0 +1,37 @@
+WITH `bfcte_1` AS (
+  SELECT
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_2` AS (
+  SELECT
+    `bfcol_1` AS `bfcol_2`,
+    `bfcol_0` AS `bfcol_3`
+  FROM `bfcte_1`
+), `bfcte_0` AS (
+  SELECT
+    `int64_too` AS `bfcol_4`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+), `bfcte_3` AS (
+  SELECT
+    `bfcte_2`.*,
+    EXISTS(
+      SELECT
+        1
+      FROM (
+        SELECT
+          `bfcol_4`
+        FROM `bfcte_0`
+        GROUP BY
+          `bfcol_4`
+      ) AS `bft_0`
+      WHERE
+        COALESCE(`bfcte_2`.`bfcol_3`, 0) = COALESCE(`bft_0`.`bfcol_4`, 0)
+        AND COALESCE(`bfcte_2`.`bfcol_3`, 1) = COALESCE(`bft_0`.`bfcol_4`, 1)
+    ) AS `bfcol_5`
+  FROM `bfcte_2`
+)
+SELECT
+  `bfcol_2` AS `rowindex`,
+  `bfcol_5` AS `int64_col`
+FROM `bfcte_3`