chore: Add support for accessing BI Engine statistics

The REST API returns BiEngineStatistics for a query which denotes if the
query was accelerated by BI Engine or not. This commit adds the
necessary function to access this information for executed queries.

Author: anmolsahoo25

google/cloud/bigquery/enums.py

@@ -28,6 +28,46 @@ class AutoRowIDs(enum.Enum):
     GENERATE_UUID = enum.auto()
 
 
+class BiEngineReasonCode(enum.Enum):
+    """Specifies reason why BI Engine did not accelerate query"""
+
+    CODE_UNSPECIFIED = enum.auto()
+    """BiEngineReason not specified."""
+
+    NO_RESERVATION = enum.auto()
+    """No reservation available for BI Engine acceleration."""
+
+    INSUFFICIENT_RESERVATION = enum.auto()
+    """Not enough memory available for BI Engine acceleration."""
+
+    UNCACHED = enum.auto()
+    """Data is not-cached and could not be accelerated by BI Engine."""
+
+    UNSUPPORTED_SQL_TEXT = enum.auto()
+    """This particular SQL text is not supported for acceleration by BI Engine."""
+
+    INPUT_TOO_LARGE = enum.auto()
+    """Input too large for acceleration by BI Engine."""
+
+    OTHER_REASON = enum.auto()
+    """Catch-all code for all other cases for partial or disabled acceleration."""
+
+    TABLE_EXCLUDED = enum.auto()
+    """One or more tables were not eligible for BI Engine acceleration."""
+
+
+class BiEngineMode(enum.Enum):
+    """Specifies which mode of BI Engine acceleration was performed"""
+
+    ACCELERATION_MODE_UNSPECIFIED = enum.auto()
+
+    DISABLED = enum.auto()
+
+    PARTIAL = enum.auto()
+
+    FULL = enum.auto()
+
+
 class Compression(object):
     """The compression type to use for exported files. The default value is
     :attr:`NONE`.

google/cloud/bigquery/job/query.py

@@ -29,6 +29,7 @@
 from google.cloud.bigquery.dataset import DatasetReference
 from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 from google.cloud.bigquery.enums import KeyResultStatementKind
+from google.cloud.bigquery.enums import BiEngineMode, BiEngineReasonCode
 from google.cloud.bigquery.external_config import ExternalConfig
 from google.cloud.bigquery import _helpers
 from google.cloud.bigquery.query import (
@@ -121,6 +122,50 @@ def _to_api_repr_table_defs(value):
     return {k: ExternalConfig.to_api_repr(v) for k, v in value.items()}
 
 
+class BiEngineReason(typing.NamedTuple):
+    """Reason for BI Engine acceleration failure
+
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#bienginereason
+    """
+
+    code: BiEngineReasonCode = BiEngineReasonCode.CODE_UNSPECIFIED
+
+    reason: str = ""
+
+    @classmethod
+    def from_api_repr(cls, reason: Dict[str, str]) -> "BiEngineReason":
+        return cls(BiEngineReasonCode[reason.get("code")], reason.get("message"))
+
+
+class BiEngineStats(typing.NamedTuple):
+    """Statistics for a BI Engine query
+
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#bienginestatistics
+    """
+
+    mode: BiEngineMode = BiEngineMode.ACCELERATION_MODE_UNSPECIFIED
+    """ Specifies which mode of BI Engine acceleration was performed (if any)
+    """
+
+    reasons: List[BiEngineReason] = []
+    """ Contains explanatory messages in case of DISABLED / PARTIAL acceleration
+    """
+
+    @classmethod
+    def from_api_repr(cls, stats: Dict[str, str]) -> "BiEngineStats":
+        biEngineMode = stats.get("biEngineMode")
+        biEngineReasons = stats.get("biEngineReasons")
+
+        mode = BiEngineMode[biEngineMode]
+
+        if biEngineReasons is None:
+            reasons = []
+        else:
+            reasons = [BiEngineReason.from_api_repr(r) for r in biEngineReasons]
+
+        return cls(mode, reasons)
+
+
 class DmlStats(typing.NamedTuple):
     """Detailed statistics for DML statements.
 
@@ -1191,6 +1236,15 @@ def dml_stats(self) -> Optional[DmlStats]:
         else:
             return DmlStats.from_api_repr(stats)
 
+    @property
+    def bi_engine_stats(self) -> Optional[BiEngineStats]:
+        stats = self._job_statistics().get("biEngineStatistics")
+
+        if stats is not None:
+            return None
+        else:
+            return BiEngineStats.from_api_repr(stats)
+
     def _blocking_poll(self, timeout=None, **kwargs):
         self._done_timeout = timeout
         self._transport_timeout = timeout

tests/unit/job/test_query_stats.py

@@ -13,6 +13,63 @@
 # limitations under the License.
 
 from .helpers import _Base
+from google.cloud.bigquery.enums import BiEngineMode, BiEngineReasonCode
+
+
+class TestBiEngineStats:
+    @staticmethod
+    def _get_target_class():
+        from google.cloud.bigquery.job.query import BiEngineStats
+
+        return BiEngineStats
+
+    def _make_one(self, *args, **kw):
+        return self._get_target_class()(*args, **kw)
+
+    def test_ctor_defaults(self):
+        bi_engine_stats = self._make_one()
+        assert bi_engine_stats.mode == BiEngineMode.ACCELERATION_MODE_UNSPECIFIED
+        assert bi_engine_stats.reasons == []
+
+    def test_from_api_repr_unspecified(self):
+        klass = self._get_target_class()
+        result = klass.from_api_repr({"biEngineMode": "ACCELERATION_MODE_UNSPECIFIED"})
+
+        assert isinstance(result, klass)
+        assert result.mode == BiEngineMode.ACCELERATION_MODE_UNSPECIFIED
+        assert result.reasons == []
+
+    def test_from_api_repr_full(self):
+        klass = self._get_target_class()
+        result = klass.from_api_repr({"biEngineMode": "FULL"})
+
+        assert isinstance(result, klass)
+        assert result.mode == BiEngineMode.FULL
+        assert result.reasons == []
+
+    def test_from_api_repr_disabled(self):
+        klass = self._get_target_class()
+        result = klass.from_api_repr(
+            {
+                "biEngineMode": "DISABLED",
+                "biEngineReasons": [
+                    {
+                        "code": "OTHER_REASON",
+                        "message": "Unable to support input table xyz due to an internal error.",
+                    }
+                ],
+            }
+        )
+
+        assert isinstance(result, klass)
+        assert result.mode == BiEngineMode.DISABLED
+
+        reason = result.reasons[0]
+        assert reason.code == BiEngineReasonCode.OTHER_REASON
+        assert (
+            reason.reason
+            == "Unable to support input table xyz due to an internal error."
+        )
 
 
 class TestDmlStats: