Add timeline and top-level slot-millis to query statistics. (googleapis#5312)

shollyman · web-flow · commit c170ee831e2c · 2018-05-09T15:56:30.000-07:00
* Add timeline and top-level slot-millis to query statistics.

* address reviewer comment:  add _int_or_none guards

* Add a system test that examines query statistics.

* Remove ratio evaluation from the query statistics system test.
diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
@@ -2200,6 +2200,14 @@ def query_plan(self):
         plan_entries = self._job_statistics().get('queryPlan', ())
         return [QueryPlanEntry.from_api_repr(entry) for entry in plan_entries]
 
+    @property
+    def timeline(self):
+        """List(TimelineEntry): Return the query execution timeline
+        from job statistics.
+        """
+        raw = self._job_statistics().get('timeline', ())
+        return [TimelineEntry.from_api_repr(entry) for entry in raw]
+
     @property
     def total_bytes_processed(self):
         """Return total bytes processed from job statistics, if present.
@@ -2274,6 +2282,11 @@ def num_dml_affected_rows(self):
             result = int(result)
         return result
 
+    @property
+    def slot_millis(self):
+        """Union[int, None]: Slot-milliseconds used by this query job."""
+        return _int_or_none(self._job_statistics().get('totalSlotMs'))
+
     @property
     def statement_type(self):
         """Return statement type from job statistics, if present.
@@ -2518,46 +2531,49 @@ def start(self):
         if self._properties.get('startMs') is None:
             return None
         return _datetime_from_microseconds(
-                self._properties.get('startMs') * 1000.0)
+                int(self._properties.get('startMs')) * 1000.0)
 
     @property
     def end(self):
         """Union[Datetime, None]: Datetime when the stage ended."""
         if self._properties.get('endMs') is None:
             return None
         return _datetime_from_microseconds(
-                self._properties.get('endMs') * 1000.0)
+                int(self._properties.get('endMs')) * 1000.0)
 
     @property
     def input_stages(self):
         """List(int): Entry IDs for stages that were inputs for this stage."""
-        return self._properties.get('inputStages', [])
+        if self._properties.get('inputStages') is None:
+            return []
+        return [_int_or_none(entry)
+                for entry in self._properties.get('inputStages')]
 
     @property
     def parallel_inputs(self):
         """Union[int, None]: Number of parallel input segments within
         the stage.
         """
-        return self._properties.get('parallelInputs')
+        return _int_or_none(self._properties.get('parallelInputs'))
 
     @property
     def completed_parallel_inputs(self):
         """Union[int, None]: Number of parallel input segments completed."""
-        return self._properties.get('completedParallelInputs')
+        return _int_or_none(self._properties.get('completedParallelInputs'))
 
     @property
     def wait_ms_avg(self):
         """Union[int, None]: Milliseconds the average worker spent waiting to
         be scheduled.
         """
-        return self._properties.get('waitMsAvg')
+        return _int_or_none(self._properties.get('waitMsAvg'))
 
     @property
     def wait_ms_max(self):
         """Union[int, None]: Milliseconds the slowest worker spent waiting to
         be scheduled.
         """
-        return self._properties.get('waitMsMax')
+        return _int_or_none(self._properties.get('waitMsMax'))
 
     @property
     def wait_ratio_avg(self):
@@ -2580,14 +2596,14 @@ def read_ms_avg(self):
         """Union[int, None]: Milliseconds the average worker spent reading
         input.
         """
-        return self._properties.get('readMsAvg')
+        return _int_or_none(self._properties.get('readMsAvg'))
 
     @property
     def read_ms_max(self):
         """Union[int, None]: Milliseconds the slowest worker spent reading
         input.
         """
-        return self._properties.get('readMsMax')
+        return _int_or_none(self._properties.get('readMsMax'))
 
     @property
     def read_ratio_avg(self):
@@ -2610,14 +2626,14 @@ def compute_ms_avg(self):
         """Union[int, None]: Milliseconds the average worker spent on CPU-bound
         processing.
         """
-        return self._properties.get('computeMsAvg')
+        return _int_or_none(self._properties.get('computeMsAvg'))
 
     @property
     def compute_ms_max(self):
         """Union[int, None]: Milliseconds the slowest worker spent on CPU-bound
         processing.
         """
-        return self._properties.get('computeMsMax')
+        return _int_or_none(self._properties.get('computeMsMax'))
 
     @property
     def compute_ratio_avg(self):
@@ -2640,14 +2656,14 @@ def write_ms_avg(self):
         """Union[int, None]: Milliseconds the average worker spent writing
         output data.
         """
-        return self._properties.get('writeMsAvg')
+        return _int_or_none(self._properties.get('writeMsAvg'))
 
     @property
     def write_ms_max(self):
         """Union[int, None]: Milliseconds the slowest worker spent writing
         output data.
         """
-        return self._properties.get('writeMsMax')
+        return _int_or_none(self._properties.get('writeMsMax'))
 
     @property
     def write_ratio_avg(self):
@@ -2668,12 +2684,12 @@ def write_ratio_max(self):
     @property
     def records_read(self):
         """Union[int, None]: Number of records read by this stage."""
-        return self._properties.get('recordsRead')
+        return _int_or_none(self._properties.get('recordsRead'))
 
     @property
     def records_written(self):
         """Union[int, None]: Number of records written by this stage."""
-        return self._properties.get('recordsWritten')
+        return _int_or_none(self._properties.get('recordsWritten'))
 
     @property
     def status(self):
@@ -2685,14 +2701,14 @@ def shuffle_output_bytes(self):
         """Union[int, None]: Number of bytes written by this stage to
         intermediate shuffle.
         """
-        return self._properties.get('shuffleOutputBytes')
+        return _int_or_none(self._properties.get('shuffleOutputBytes'))
 
     @property
     def shuffle_output_bytes_spilled(self):
         """Union[int, None]: Number of bytes written by this stage to
         intermediate shuffle and spilled to disk.
         """
-        return self._properties.get('shuffleOutputBytesSpilled')
+        return _int_or_none(self._properties.get('shuffleOutputBytesSpilled'))
 
     @property
     def steps(self):
@@ -2703,6 +2719,66 @@ def steps(self):
                 for step in self._properties.get('steps', [])]
 
 
+class TimelineEntry(object):
+    """TimelineEntry represents progress of a query job at a particular
+    point in time.
+
+    See
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs
+    for the underlying API representation within query statistics.
+
+    """
+
+    def __init__(self):
+        self._properties = {}
+
+    @classmethod
+    def from_api_repr(cls, resource):
+        """Factory: construct instance from the JSON repr.
+
+        Args:
+            resource(Dict[str: object]):
+                QueryTimelineSample representation returned from API
+
+        Returns:
+            google.cloud.bigquery.TimelineEntry:
+                Timeline sample parsed from ``resource``
+        """
+        entry = cls()
+        entry._properties = resource
+        return entry
+
+    @property
+    def elapsed_ms(self):
+        """Union[int, None]: Milliseconds elapsed since start of query
+        execution."""
+        return _int_or_none(self._properties.get('elapsedMs'))
+
+    @property
+    def active_units(self):
+        """Union[int, None]: Current number of input units being processed
+        by workers, reported as largest value since the last sample."""
+        return _int_or_none(self._properties.get('activeUnits'))
+
+    @property
+    def pending_units(self):
+        """Union[int, None]: Current number of input units remaining for
+        query stages active at this sample time."""
+        return _int_or_none(self._properties.get('pendingUnits'))
+
+    @property
+    def completed_units(self):
+        """Union[int, None]: Current number of input units completed by
+        this query."""
+        return _int_or_none(self._properties.get('completedUnits'))
+
+    @property
+    def slot_millis(self):
+        """Union[int, None]: Cumulative slot-milliseconds consumed by
+        this query."""
+        return _int_or_none(self._properties.get('totalSlotMs'))
+
+
 class UnknownJob(_AsyncJob):
     """A job whose type cannot be determined."""
 
diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py
@@ -1087,6 +1087,82 @@ def test_query_w_timeout(self):
             # 1 second is much too short for this query.
             query_job.result(timeout=1)
 
+    def test_query_statistics(self):
+        """
+        A system test to exercise some of the extended query statistics.
+
+        Note:  We construct a query that should need at least three stages by
+        specifying a JOIN query.  Exact plan and stats are effectively
+        non-deterministic, so we're largely interested in confirming values
+        are present.
+        """
+
+        job_config = bigquery.QueryJobConfig()
+        job_config.use_query_cache = False
+
+        query_job = Config.CLIENT.query(
+            """
+            SELECT
+              COUNT(1)
+            FROM
+            (
+              SELECT
+                year,
+                wban_number
+              FROM `bigquery-public-data.samples.gsod`
+              LIMIT 1000
+            ) lside
+            INNER JOIN
+            (
+              SELECT
+                year,
+                state
+              FROM `bigquery-public-data.samples.natality`
+              LIMIT 1000
+            ) rside
+            ON
+            lside.year = rside.year
+            """,
+            location='US',
+            job_config=job_config)
+
+        # run the job to completion
+        query_job.result()
+
+        # Assert top-level stats
+        self.assertFalse(query_job.cache_hit)
+        self.assertIsNotNone(query_job.destination)
+        self.assertTrue(query_job.done)
+        self.assertFalse(query_job.dry_run)
+        self.assertIsNone(query_job.num_dml_affected_rows)
+        self.assertEqual(query_job.priority, 'INTERACTIVE')
+        self.assertGreater(query_job.total_bytes_billed, 1)
+        self.assertGreater(query_job.total_bytes_processed, 1)
+        self.assertEqual(query_job.statement_type, 'SELECT')
+        self.assertGreater(query_job.slot_millis, 1)
+
+        # Make assertions on the shape of the query plan.
+        plan = query_job.query_plan
+        self.assertGreaterEqual(len(plan), 3)
+        first_stage = plan[0]
+        self.assertIsNotNone(first_stage.start)
+        self.assertIsNotNone(first_stage.end)
+        self.assertIsNotNone(first_stage.entry_id)
+        self.assertIsNotNone(first_stage.name)
+        self.assertGreater(first_stage.parallel_inputs, 0)
+        self.assertGreater(first_stage.completed_parallel_inputs, 0)
+        self.assertGreater(first_stage.shuffle_output_bytes, 0)
+        self.assertEqual(first_stage.status, 'COMPLETE')
+
+        # Query plan is a digraph.  Ensure it has inter-stage links,
+        # but not every stage has inputs.
+        stages_with_inputs = 0
+        for entry in plan:
+            if len(entry.input_stages) > 0:
+                stages_with_inputs = stages_with_inputs + 1
+        self.assertGreater(stages_with_inputs, 0)
+        self.assertGreater(len(plan), stages_with_inputs)
+
     def test_dbapi_w_standard_sql_types(self):
         examples = self._generate_standard_sql_types_examples()
         for example in examples:
diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
