feat(connect): printSchema

Cargo.toml

@@ -190,6 +190,7 @@ chrono = "0.4.38"
 chrono-tz = "0.10.0"
 comfy-table = "7.1.1"
 common-daft-config = {path = "src/common/daft-config"}
+common-display = {path = "src/common/display", default-features = false}
 common-error = {path = "src/common/error", default-features = false}
 common-file-formats = {path = "src/common/file-formats"}
 common-runtime = {path = "src/common/runtime", default-features = false}

src/daft-connect/Cargo.toml

@@ -2,6 +2,7 @@
 arrow2 = {workspace = true, features = ["io_json_integration"]}
 async-stream = "0.3.6"
 common-daft-config = {workspace = true}
+common-display = {workspace = true}
 common-file-formats = {workspace = true}
 daft-core = {workspace = true}
 daft-dsl = {workspace = true}

src/daft-connect/src/lib.rs

@@ -6,6 +6,8 @@
 #![feature(stmt_expr_attributes)]
 #![feature(try_trait_v2_residual)]
 
+use common_display::{tree::TreeDisplay, DisplayLevel};
+use daft_micropartition::partitioning::InMemoryPartitionSetCache;
 use dashmap::DashMap;
 use eyre::Context;
 #[cfg(feature = "python")]
@@ -22,10 +24,10 @@ use spark_connect::{
     ReleaseExecuteResponse, ReleaseSessionRequest, ReleaseSessionResponse,
 };
 use tonic::{transport::Server, Request, Response, Status};
-use tracing::info;
+use tracing::{info, warn};
 use uuid::Uuid;
 
-use crate::session::Session;
+use crate::{session::Session, translation::SparkAnalyzer};
 
 mod config;
 mod err;
@@ -303,7 +305,7 @@ impl SparkConnectService for DaftSparkConnectService {
                     return Err(Status::invalid_argument("op_type is required to be root"));
                 };
 
-                let result = match translation::relation_to_schema(relation).await {
+                let result = match translation::relation_to_spark_schema(relation).await {
                     Ok(schema) => schema,
                     Err(e) => {
                         return invalid_argument_err!(
@@ -346,6 +348,45 @@ impl SparkConnectService for DaftSparkConnectService {
 
                 Ok(Response::new(response))
             }
+            Analyze::TreeString(TreeString { plan, level }) => {
+                let Some(plan) = plan else {
+                    return invalid_argument_err!("plan is required");
+                };
+
+                let Some(op_type) = plan.op_type else {
+                    return invalid_argument_err!("op_type is required");
+                };
+
+                let OpType::Root(input) = op_type else {
+                    return invalid_argument_err!("op_type must be Root");
+                };
+
+                if let Some(common) = &input.common {
+                    if common.origin.is_some() {
+                        warn!("Ignoring common metadata for relation: {common:?}; not yet implemented");
+                    }
+                }
+
+                // We're just checking the schema here, so we don't need to use a persistent cache as it won't be used
+                let pset = InMemoryPartitionSetCache::empty();
+                let translator = SparkAnalyzer::new(&pset);
+                let plan = Box::pin(translator.to_logical_plan(input))
+                    .await
+                    .unwrap()
+                    .build();
+
+                let s = plan.display_as(DisplayLevel::Default);
+
+                let response = AnalyzePlanResponse {
+                    session_id,
+                    server_side_session_id: String::new(),
+                    result: Some(spark_connect::analyze_plan_response::Result::TreeString(
+                        analyze_plan_response::TreeString { tree_string: s },
+                    )),
+                };
+
+                Ok(Response::new(response))
+            }
             other => unimplemented_err!("Analyze plan operation is not yet implemented: {other:?}"),
         }
     }

src/daft-connect/src/translation.rs

@@ -10,4 +10,4 @@ pub use datatype::{to_daft_datatype, to_spark_datatype};
 pub use expr::to_daft_expr;
 pub use literal::to_daft_literal;
 pub use logical_plan::SparkAnalyzer;
-pub use schema::relation_to_schema;
+pub use schema::{relation_to_daft_schema, relation_to_spark_schema};

src/daft-connect/src/translation/schema.rs

@@ -1,4 +1,5 @@
 use daft_micropartition::partitioning::InMemoryPartitionSetCache;
+use daft_schema::schema::{Schema, SchemaRef};
 use spark_connect::{
     data_type::{Kind, Struct, StructField},
     DataType, Relation,
@@ -9,19 +10,8 @@ use super::SparkAnalyzer;
 use crate::translation::to_spark_datatype;
 
 #[tracing::instrument(skip_all)]
-pub async fn relation_to_schema(input: Relation) -> eyre::Result<DataType> {
-    if let Some(common) = &input.common {
-        if common.origin.is_some() {
-            warn!("Ignoring common metadata for relation: {common:?}; not yet implemented");
-        }
-    }
-
-    // We're just checking the schema here, so we don't need to use a persistent cache as it won't be used
-    let pset = InMemoryPartitionSetCache::empty();
-    let translator = SparkAnalyzer::new(&pset);
-    let plan = Box::pin(translator.to_logical_plan(input)).await?;
-
-    let result = plan.schema();
+pub async fn relation_to_spark_schema(input: Relation) -> eyre::Result<DataType> {
+    let result = relation_to_daft_schema(input).await?;
 
     let fields: eyre::Result<Vec<StructField>> = result
         .fields
@@ -44,3 +34,21 @@ pub async fn relation_to_schema(input: Relation) -> eyre::Result<DataType> {
         })),
     })
 }
+
+#[tracing::instrument(skip_all)]
+pub async fn relation_to_daft_schema(input: Relation) -> eyre::Result<SchemaRef> {
+    if let Some(common) = &input.common {
+        if common.origin.is_some() {
+            warn!("Ignoring common metadata for relation: {common:?}; not yet implemented");
+        }
+    }
+
+    // We're just checking the schema here, so we don't need to use a persistent cache as it won't be used
+    let pset = InMemoryPartitionSetCache::empty();
+    let translator = SparkAnalyzer::new(&pset);
+    let plan = Box::pin(translator.to_logical_plan(input)).await?;
+
+    let result = plan.schema();
+
+    Ok(result)
+}

tests/connect/test_print_schema.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+
+def test_print_schema(spark_session: object, capsys: object) -> None:
+    df = spark_session.range(10)
+    df.printSchema()
+
+    captured = capsys.readouterr()
+    expected = (
+        "root\n"
+        " |-- id: long (nullable = true)\n"
+    )
+    assert captured.out == expected