Skip to content

Commit

Permalink
DB Connect Progress: Make sure we always end up at 100% (#1363)
Browse files Browse the repository at this point in the history
## Changes
DB Connect Progress: Make sure we always end up at 100%
## Tests
<!-- How is this tested? -->
  • Loading branch information
fjakobs authored Sep 18, 2024
1 parent a9f79d2 commit e51e62b
Showing 1 changed file with 14 additions and 6 deletions.
20 changes: 14 additions & 6 deletions packages/databricks-vscode/resources/python/00-databricks-init.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,14 @@
import json
from typing import Any, Union, List
import os
import sys
import time
import shlex
import warnings
import tempfile

# prevent sum from pyskaprk.sql.functions from shadowing the builtin sum
builtinSum = sys.modules['builtins'].sum

def logError(function_name: str, e: Union[str, Exception]):
import sys
Expand Down Expand Up @@ -403,14 +406,20 @@ def init_ui(self):
def update_ticks(
self,
stages,
inflight_tasks: int
inflight_tasks: int,
done: bool
) -> None:
total_tasks = sum(map(lambda x: x.num_tasks, stages))
completed_tasks = sum(map(lambda x: x.num_completed_tasks, stages))
total_tasks = builtinSum(map(lambda x: x.num_tasks, stages))
completed_tasks = builtinSum(map(lambda x: x.num_completed_tasks, stages))
if total_tasks > 0:
self._ticks = total_tasks
self._tick = completed_tasks
self._bytes_read = sum(map(lambda x: x.num_bytes_read, stages))
self._bytes_read = builtinSum(map(lambda x: x.num_bytes_read, stages))

if done:
self._tick = self._ticks
self._running = 0

if self._tick is not None and self._tick >= 0:
self.output()
self._running = inflight_tasks
Expand All @@ -432,7 +441,6 @@ def _bytes_to_string(size: int) -> str:
i += 1
result = float(size) / Progress.SI_BYTE_SIZES[i]
return f"{result:.1f} {Progress.SI_BYTE_SUFFIXES[i]}"


class ProgressHandler:
def __init__(self):
Expand All @@ -454,7 +462,7 @@ def __call__(self,
self.op_id = operation_id
self.reset()

self.p.update_ticks(stages, inflight_tasks)
self.p.update_ticks(stages, inflight_tasks, done)

spark.clearProgressHandlers()
spark.registerProgressHandler(ProgressHandler())
Expand Down

0 comments on commit e51e62b

Please sign in to comment.