feat: Broadcast plan over port 3238

daft/context.py

@@ -125,6 +125,10 @@
 
     _runner: Runner | None = None
 
+    _enable_broadcast: bool = False
+    _broadcast_addr: str = "127.0.0.1"
+    _broadcast_port: int = 3238
+
     _instance: ClassVar[DaftContext | None] = None
     _lock: ClassVar[threading.Lock] = threading.Lock()
 
@@ -197,6 +201,31 @@
     return _DaftContext
 
 
+def broadcast_metrics(
+    enable: bool = True,
+    addr: str = "127.0.0.1",
+    port: int = 3238,
+) -> DaftContext:
+    """Enable (or disable) the broadcasting of Daft metrics over the given address.
+
+    This enables the Daft Dashboard to collect and display the data which is sent to it as Daft executes queries.
+
+    Args:
+        enable: Whether or not to enable broadcasting Daft metrics.
+        addr: The IP address on which to broadcast the metrics.
+        port: The port on which to broadcast the metrics.
+
+    Returns:
+        DaftContext: Daft context after enabling or disabling metrics broadcasting.
+    """
+    ctx = get_context()
+    with ctx._lock:
+        ctx._enable_broadcast = enable
+        ctx._broadcast_addr = addr
+        ctx._broadcast_port = port
+        return ctx
+
+
 def set_runner_ray(
     address: str | None = None,
     noop_if_initialized: bool = False,

daft/dataframe/dataframe.py

@@ -158,6 +158,28 @@
         else:
             return self._result_cache.value
 
+    def explain_broadcast(self):
+        """Broadcast the mermaid-formatted plan on the given port (assuming metrics-broadcasting is enabled)."""
+        import socket
+
+        from daft.dataframe.display import MermaidFormatter
+
+        ctx = get_context()
+        with ctx._lock:
+            if not ctx._enable_broadcast:
+                return
+
+            addr = ctx._broadcast_addr
+            port = ctx._broadcast_port
+
+        is_cached = self._result_cache is not None
+        instance = MermaidFormatter(builder=self.__builder, show_all=True, simple=False, is_cached=is_cached)
+        text: str = instance._repr_markdown_()
+
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.connect((addr, port))
+            sock.sendall(text.encode("utf-8"))
+
     @DataframePublicAPI
     def explain(
         self, show_all: bool = False, format: str = "ascii", simple: bool = False, file: Optional[io.IOBase] = None
@@ -177,6 +199,7 @@
             file (Optional[io.IOBase]): Location to print the output to, or defaults to None which defaults to the default location for
                 print (in Python, that should be sys.stdout)
         """
+        self.explain_broadcast()
         is_cached = self._result_cache is not None
         if format == "mermaid":
             from daft.dataframe.display import MermaidFormatter
@@ -2818,6 +2841,7 @@
             DataFrame: DataFrame with materialized results.
         """
         self._materialize_results()
+        self.explain_broadcast()
 
         assert self._result is not None
         dataframe_len = len(self._result)
@@ -2889,6 +2913,7 @@
             n: number of rows to show. Defaults to 8.
         """
         dataframe_display = self._construct_show_display(n)
+        self.explain_broadcast()
         try:
             from IPython.display import display