From fb62664286700516cc4ffc47c7937439c41c8435 Mon Sep 17 00:00:00 2001
From: Sammy Sidhu <samster25@users.noreply.github.com>
Date: Wed, 5 Jun 2024 11:59:37 -0700
Subject: [PATCH] Introduce terminal wrap around when explaining plans

---
 Cargo.lock                   | 35 +++++++++++++++++++++++++++++++++++
 src/daft-plan/Cargo.toml     |  2 ++
 src/daft-plan/src/display.rs | 31 +++++++++++++++++++++++++------
 3 files changed, 62 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ba1e970fed..73adf97a78 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1523,6 +1523,8 @@ dependencies = [
  "serde",
  "serde_json",
  "snafu",
+ "terminal_size",
+ "textwrap",
 ]
 
 [[package]]
@@ -3824,6 +3826,12 @@ version = "1.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "942b4a808e05215192e39f4ab80813e599068285906cc91aa64f923db842bd5a"
 
+[[package]]
+name = "smawk"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b7c388c1b5e93756d0c740965c41e8822f866621d41acbdf6336a6a168f8840c"
+
 [[package]]
 name = "snafu"
 version = "0.7.5"
@@ -4031,6 +4039,27 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "terminal_size"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "21bebf2b7c9e0a515f6e0f8c51dc0f8e4696391e6f1ff30379559f8365fb0df7"
+dependencies = [
+ "rustix",
+ "windows-sys 0.48.0",
+]
+
+[[package]]
+name = "textwrap"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23d434d3f8967a09480fb04132ebe0a3e088c173e6d0ee7897abbdf4eab0f8b9"
+dependencies = [
+ "smawk",
+ "unicode-linebreak",
+ "unicode-width",
+]
+
 [[package]]
 name = "thiserror"
 version = "1.0.49"
@@ -4325,6 +4354,12 @@ version = "1.0.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
 
+[[package]]
+name = "unicode-linebreak"
+version = "0.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3b09c83c3c29d37506a3e260c08c03743a6bb66a9cd432c6934ab501a190571f"
+
 [[package]]
 name = "unicode-normalization"
 version = "0.1.22"
diff --git a/src/daft-plan/Cargo.toml b/src/daft-plan/Cargo.toml
index ccfbbf1bd7..67880343c6 100644
--- a/src/daft-plan/Cargo.toml
+++ b/src/daft-plan/Cargo.toml
@@ -18,6 +18,8 @@ pyo3-log = {workspace = true, optional = true}
 serde = {workspace = true, features = ["rc"]}
 serde_json = {workspace = true}
 snafu = {workspace = true}
+terminal_size = {version = "0.3.0"}
+textwrap = {version = "0.16.1"}
 
 [dev-dependencies]
 rstest = {workspace = true}
diff --git a/src/daft-plan/src/display.rs b/src/daft-plan/src/display.rs
index 03f9d8d28a..e896280984 100644
--- a/src/daft-plan/src/display.rs
+++ b/src/daft-plan/src/display.rs
@@ -58,13 +58,32 @@ pub(crate) trait TreeDisplay {
         } else {
             self.get_multiline_representation()
         };
-        for (i, val) in lines.iter().enumerate() {
-            self.fmt_depth(depth, s)?;
-            match i {
-                0 => write!(s, "* ")?,
-                _ => write!(s, "|   ")?,
+        use terminal_size::{terminal_size, Width};
+        let size = terminal_size();
+        let term_width = if let Some((Width(w), _)) = size {
+            w as usize
+        } else {
+            88usize
+        };
+
+        let mut counter = 0;
+        for val in lines.iter() {
+            let base_characters = depth * 2;
+            let expected_chars = (term_width - base_characters - 8).max(8);
+            let sublines = textwrap::wrap(val, expected_chars);
+
+            for (i, sb) in sublines.iter().enumerate() {
+                self.fmt_depth(depth, s)?;
+                match counter {
+                    0 => write!(s, "* ")?,
+                    _ => write!(s, "|   ")?,
+                }
+                counter += 1;
+                match i {
+                    0 => writeln!(s, "{sb}")?,
+                    _ => writeln!(s, "  {sb}")?,
+                }
             }
-            writeln!(s, "{val}")?;
         }
 
         // Recursively handle children.