Add support for accessing BI Engine statistics of query

anmolsahoo25 · anmolsahoo25 · commit ef070537a0a3 · 2022-02-18T02:42:53.000-05:00
The REST API return BiEngineStatistics 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.
diff --git a/google/cloud/bigquery/enums.py b/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`.
diff --git a/google/cloud/bigquery/job/query.py b/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
diff --git a/tests/unit/job/test_query_stats.py b/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:
