Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

DB Connect Progress: Make sure we always end up at 100% #1363

Merged
merged 6 commits into from
Sep 18, 2024
Merged

DB Connect Progress: Make sure we always end up at 100% #1363

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
31ff699
Inline code from pyspark
fjakobs Sep 16, 2024
ea9d999
Add option to disable progress through an environment variable
fjakobs Sep 16, 2024
914a978
Expose db connect progress as a settings
fjakobs Sep 16, 2024
44200c7
Reload ipy kernel on settings change
fjakobs Sep 16, 2024
37d1458
Also show progress for running python files
fjakobs Sep 17, 2024
afde666
Make sure we always end up at 100%
fjakobs Sep 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 14 additions & 6 deletions packages/databricks-vscode/resources/python/00-databricks-init.py
View file
Open in desktop
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))
fjakobs marked this conversation as resolved.
Show resolved Hide resolved
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
Loading