chore: sync changes from internal repo (#15)

feat: support `DataFrame.isin` with list and dict inputs
test: move flaky `reset_session` test to unit tests
chore: don't run redundant tests in nightly build
feat: support `DataFrame`-`DataFrame` binary operations
feat: support `Series.map`
feat: support `Index.is_monotonic`
docs: update remote function notebook with read_gbq_function usage
feat: use default session and connection in `ml.llm` and `ml.imported`
chore: disable broken stack tests
feat: support `pow()` and power operator in `DataFrame` and `Series`
feat: support for `np.add`, `np.subtract`, `np.multiply`, `np.divide`, `np.power`
perf: use `row_number()` filter for `head` / `tail`
feat: support `bigframes.pandas.merge()`
fix: make `X_train` argument names consistent across methods
chore: refactor ml core
feat: add `Series.dropna` and `DataFrame.fillna`
chore: fix gcs notebooks upload in 'nightly' build
chore: fix Kokoro build files to support GitHub directories
chore: fix unit test to not require authentication

Author: tswast

.kokoro/build.sh

@@ -15,8 +15,14 @@
 
 set -eo pipefail
 
+if [[ -z "${KOKORO_GOB_COMMIT}" ]]; then
+    PROJECT_SCM="github"
+else
+    PROJECT_SCM="git"
+fi
+
 if [[ -z "${PROJECT_ROOT:-}" ]]; then
-    PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/git/bigframes"
+    PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}/bigframes"
 fi
 
 cd "${PROJECT_ROOT}"

.kokoro/continuous/nightly.cfg

@@ -1,8 +1,3 @@
 # Format: //devtools/kokoro/config/proto/build.proto
 
-env_vars: {
-    key: "NOX_SESSION"
-    value: "unit system cover lint lint_setup_py mypy format docs e2e notebook"
-}
-
 build_file: "bigframes/.kokoro/release-nightly.sh"

.kokoro/release-nightly.sh

@@ -34,8 +34,14 @@ while [ $# -gt 0 ] ; do
   shift 1;
 done
 
+if [[ -z "${KOKORO_GOB_COMMIT}" ]]; then
+    PROJECT_SCM="github"
+else
+    PROJECT_SCM="git"
+fi
+
 if [ -z "${PROJECT_ROOT:-}" ]; then
-    PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/git/bigframes"
+    PROJECT_ROOT="${KOKORO_ARTIFACTS_DIR}/${PROJECT_SCM}/bigframes"
 fi
 
 # Move into the package, build the distribution and upload to shared bucket.
@@ -57,16 +63,6 @@ export PYTHONUNBUFFERED=1
 # Install dependencies, as the following steps depend on it
 python3.10 -m pip install -e .[all]
 
-# If NOX_SESSION is set, it only runs the specified session,
-# otherwise run all the sessions.
-if ! [ ${DRY_RUN} ]; then
-    if [ -n "${NOX_SESSION:-}" ]; then
-        python3.10 -m nox -s ${NOX_SESSION:-}
-    else
-        python3.10 -m nox
-    fi
-fi
-
 # Generate third party notices and include it in the licenses in setup.cfg
 # TODO(shobs): Don't include it in the package once vertex colab can pick it
 # from elsewhere
@@ -138,15 +134,8 @@ if ! [ ${DRY_RUN} ]; then
       gsutil cp -v dist/* ${gcs_path}
       gsutil cp -v LICENSE ${gcs_path}
       gsutil cp -v ${THIRD_PARTY_NOTICES_FILE} ${gcs_path}
-      gsutil -m cp -v "notebooks/00 - Summary.ipynb" \
-                      "notebooks/01 - Getting Started.ipynb" \
-                      "notebooks/02 - DataFrame.ipynb" \
-                      "notebooks/03 - Using ML - ML fundamentals.ipynb" \
-                      "notebooks/04 - Using ML - SKLearn linear regression.ipynb" \
-                      "notebooks/05 - Using ML - Easy linear regression.ipynb" \
-                      "notebooks/06 - Using ML - Large Language Models.ipynb" \
-                      "notebooks/50 - Remote Function.ipynb" \
-                      ${gcs_path}notebooks/
+      gsutil -m cp -r -v "notebooks/" ${gcs_path}notebooks/
+
     done
 
     # publish API coverage information to BigQuery

OWNERS

@@ -5,6 +5,7 @@ henryjsolberg@google.com
 hormati@google.com
 huanc@google.com
 jiaxun@google.com
+kemppeterson@google.com
 shobs@google.com
 swast@google.com
 tbergeron@google.com

bigframes/_config/bigquery_options.py

@@ -35,13 +35,13 @@ def __init__(
         credentials: Optional[google.auth.credentials.Credentials] = None,
         project: Optional[str] = None,
         location: Optional[str] = None,
-        remote_udf_connection: Optional[str] = None,
+        bq_connection: Optional[str] = None,
         use_regional_endpoints: bool = False,
     ):
         self._credentials = credentials
         self._project = project
         self._location = location
-        self._remote_udf_connection = remote_udf_connection
+        self._bq_connection = bq_connection
         self._use_regional_endpoints = use_regional_endpoints
         self._session_started = False
 
@@ -82,23 +82,21 @@ def project(self, value: Optional[str]):
         self._project = value
 
     @property
-    def remote_udf_connection(self) -> Optional[str]:
-        """Name of the BigQuery connection to use for remote functions.
+    def bq_connection(self) -> Optional[str]:
+        """Name of the BigQuery connection to use.
 
         You should either have the connection already created in the
         <code>location</code> you have chosen, or you should have the Project IAM
         Admin role to enable the service to create the connection for you if you
         need it.
         """
-        return self._remote_udf_connection
+        return self._bq_connection
 
-    @remote_udf_connection.setter
-    def remote_udf_connection(self, value: Optional[str]):
-        if self._session_started and self._remote_udf_connection != value:
-            raise ValueError(
-                SESSION_STARTED_MESSAGE.format(attribute="remote_udf_connection")
-            )
-        self._remote_udf_connection = value
+    @bq_connection.setter
+    def bq_connection(self, value: Optional[str]):
+        if self._session_started and self._bq_connection != value:
+            raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="bq_connection"))
+        self._bq_connection = value
 
     @property
     def use_regional_endpoints(self) -> bool:

bigframes/clients.py

@@ -0,0 +1,163 @@
+# Copyright 2023 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.
+
+"""BigQuery DataFrame clients to interact with other cloud resources"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Optional
+
+import google.api_core.exceptions
+from google.cloud import bigquery_connection_v1, resourcemanager_v3
+from google.iam.v1 import iam_policy_pb2, policy_pb2
+
+logging.basicConfig(
+    level=logging.INFO, format="[%(levelname)s][%(asctime)s][%(name)s] %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+class BqConnectionManager:
+    """Manager to handle operations with BQ connections."""
+
+    # Wait time (in seconds) for an IAM binding to take effect after creation
+    _IAM_WAIT_SECONDS = 120
+
+    def __init__(
+        self,
+        bq_connection_client: bigquery_connection_v1.ConnectionServiceClient,
+        cloud_resource_manager_client: resourcemanager_v3.ProjectsClient,
+    ):
+        self._bq_connection_client = bq_connection_client
+        self._cloud_resource_manager_client = cloud_resource_manager_client
+
+    def create_bq_connection(
+        self, project_id: str, location: str, connection_id: str, iam_role: str
+    ):
+        """Create the BQ connection if not exist. In addition, try to add the IAM role to the connection to ensure required permissions.
+
+        Args:
+            project_id:
+                ID of the project.
+            location:
+                Location of the connection.
+            connection_id:
+                ID of the connection.
+            iam_role:
+                str of the IAM role that the service account of the created connection needs to aquire. E.g. 'run.invoker', 'aiplatform.user'
+        """
+        # TODO(shobs): The below command to enable BigQuery Connection API needs
+        # to be automated. Disabling for now since most target users would not
+        # have the privilege to enable API in a project.
+        # log("Making sure BigQuery Connection API is enabled")
+        # if os.system("gcloud services enable bigqueryconnection.googleapis.com"):
+        #    raise ValueError("Failed to enable BigQuery Connection API")
+        # If the intended connection does not exist then create it
+        service_account_id = self._get_service_account_if_connection_exists(
+            project_id, location, connection_id
+        )
+        if service_account_id:
+            logger.info(
+                f"Connector {project_id}.{location}.{connection_id} already exists"
+            )
+        else:
+            connection_name, service_account_id = self._create_bq_connection(
+                project_id, location, connection_id
+            )
+            logger.info(
+                f"Created BQ connection {connection_name} with service account id: {service_account_id}"
+            )
+        # Ensure IAM role on the BQ connection
+        # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#grant_permission_on_function
+        self._ensure_iam_binding(project_id, service_account_id, iam_role)
+
+    # Introduce retries to accommodate transient errors like etag mismatch,
+    # which can be caused by concurrent operation on the same resource, and
+    # manifests with message like:
+    # google.api_core.exceptions.Aborted: 409 There were concurrent policy
+    # changes. Please retry the whole read-modify-write with exponential
+    # backoff. The request's ETag '\007\006\003,\264\304\337\272' did not match
+    # the current policy's ETag '\007\006\003,\3750&\363'.
+    @google.api_core.retry.Retry(
+        predicate=google.api_core.retry.if_exception_type(
+            google.api_core.exceptions.Aborted
+        ),
+        initial=10,
+        maximum=20,
+        multiplier=2,
+        timeout=60,
+    )
+    def _ensure_iam_binding(
+        self, project_id: str, service_account_id: str, iam_role: str
+    ):
+        """Ensure necessary IAM role is configured on a service account."""
+        project = f"projects/{project_id}"
+        service_account = f"serviceAccount:{service_account_id}"
+        role = f"roles/{iam_role}"
+        request = iam_policy_pb2.GetIamPolicyRequest(resource=project)
+        policy = self._cloud_resource_manager_client.get_iam_policy(request=request)
+
+        # Check if the binding already exists, and if does, do nothing more
+        for binding in policy.bindings:
+            if binding.role == role:
+                if service_account in binding.members:
+                    return
+
+        # Create a new binding
+        new_binding = policy_pb2.Binding(role=role, members=[service_account])
+        policy.bindings.append(new_binding)
+        request = iam_policy_pb2.SetIamPolicyRequest(resource=project, policy=policy)
+        self._cloud_resource_manager_client.set_iam_policy(request=request)
+
+        # We would wait for the IAM policy change to take effect
+        # https://cloud.google.com/iam/docs/access-change-propagation
+        logger.info(
+            f"Waiting {self._IAM_WAIT_SECONDS} seconds for IAM to take effect.."
+        )
+        time.sleep(self._IAM_WAIT_SECONDS)
+
+    def _create_bq_connection(self, project_id: str, location: str, connection_id: str):
+        """Create the BigQuery Connection and returns corresponding service account id."""
+        client = self._bq_connection_client
+        connection = bigquery_connection_v1.Connection(
+            cloud_resource=bigquery_connection_v1.CloudResourceProperties()
+        )
+        request = bigquery_connection_v1.CreateConnectionRequest(
+            parent=client.common_location_path(project_id, location),
+            connection_id=connection_id,
+            connection=connection,
+        )
+        connection = client.create_connection(request)
+        return connection.name, connection.cloud_resource.service_account_id
+
+    def _get_service_account_if_connection_exists(
+        self, project_id: str, location: str, connection_id: str
+    ) -> Optional[str]:
+        """Check if the BigQuery Connection exists."""
+        client = self._bq_connection_client
+        request = bigquery_connection_v1.GetConnectionRequest(
+            name=client.connection_path(project_id, location, connection_id)
+        )
+
+        service_account = None
+        try:
+            service_account = client.get_connection(
+                request=request
+            ).cloud_resource.service_account_id
+        except google.api_core.exceptions.NotFound:
+            pass
+
+        return service_account

bigframes/core/__init__.py

@@ -1021,20 +1021,6 @@ def slice(
         if not step:
             step = 1
 
-        # Special cases for head() and tail(), where we don't need to project
-        # offsets. LIMIT clause is much more efficient in BigQuery than a
-        # filter on row_number().
-        if (
-            (start is None or start == 0)
-            and step == 1
-            and stop is not None
-            and stop > 0
-        ):
-            return self.apply_limit(stop)
-
-        if start is not None and start < 0 and step == 1 and stop is None:
-            return self.reversed().apply_limit(abs(start)).reversed()
-
         expr_with_offsets = self.project_offsets()
 
         # start with True and reduce with start, stop, and step conditions