add support for Expand nodes

Author: menishmueli

spark-plugin/example_3_5_1/src/main/scala/io/dataflint/example/ExpandExample.scala

@@ -0,0 +1,32 @@
+package io.dataflint.example
+
+import org.apache.spark.sql.{SparkSession, functions => F}
+
+object ExpandExample extends App {
+  val spark = SparkSession
+    .builder()
+    .appName("Expand ETL")
+    .config("spark.plugins", "io.dataflint.spark.SparkDataflintPlugin")
+    .config("spark.ui.port", "10000")
+    .master("local[*]")
+    .getOrCreate()
+
+  val salesFilesLocation = sys.env.getOrElse("SALES_FILES_LOCATION", throw new Exception("SALES_FILES_LOCATION env var not set"))
+
+  val df = spark.read.load(salesFilesLocation)
+
+  // Query with multiple count distincts
+  val dfCounts = df.select(
+    F.countDistinct(F.col("ss_customer_sk")).alias("distinct_customers"),
+    F.countDistinct(F.col("ss_item_sk")).alias("distinct_items"),
+    F.countDistinct(F.col("ss_store_sk")).alias("distinct_stores"),
+    F.countDistinct(F.col("ss_promo_sk")).alias("distinct_promotions")
+  )
+
+  dfCounts.show()
+
+  scala.io.StdIn.readLine("job ended, press any key to continue..")
+
+  spark.stop()
+}
+

spark-ui/src/components/SqlFlow/StageNode.tsx

@@ -594,6 +594,29 @@ export const StageNode: FC<{
           );
         }
         break;
+      case "Expand":
+        if (
+          parsedPlan.plan.fields !== undefined &&
+          parsedPlan.plan.fields.length > 0
+        ) {
+          addTruncatedCodeTooltipMultiline(
+            dataTable,
+            "Expanded Fields",
+            parsedPlan.plan.fields,
+          );
+        }
+        if (parsedPlan.plan.idField !== undefined) {
+          addTruncatedSmallTooltip(
+            dataTable,
+            "ID Field",
+            parsedPlan.plan.idField,
+            25,
+            true,
+            false,
+            false,
+          );
+        }
+        break;
     }
   }
 

spark-ui/src/interfaces/AppStore.ts

@@ -255,6 +255,11 @@ export type ParsedGeneratePlan = {
   selectedFields: string[];
 };
 
+export type ParsedExpandPlan = {
+  fields: string[];
+  idField: string;
+};
+
 export type ParsedNodePlan =
   | { type: "HashAggregate"; plan: ParsedHashAggregatePlan }
   | { type: "TakeOrderedAndProject"; plan: ParsedTakeOrderedAndProjectPlan }
@@ -269,7 +274,8 @@ export type ParsedNodePlan =
   | { type: "Window"; plan: ParsedWindowPlan }
   | { type: "Coalesce"; plan: ParsedCoalescePlan }
   | { type: "BatchEvalPython"; plan: ParsedBatchEvalPythonPlan }
-  | { type: "Generate"; plan: ParsedGeneratePlan };
+  | { type: "Generate"; plan: ParsedGeneratePlan }
+  | { type: "Expand"; plan: ParsedExpandPlan };
 
 
 export interface ExchangeMetrics {

spark-ui/src/reducers/PlanParsers/ExpandParser.spec.ts

@@ -0,0 +1,41 @@
+import { parseExpand } from "./ExpandParser";
+
+describe("parseExpand", () => {
+    test("should parse expand plan correctly", () => {
+        const input =
+            "Expand [[ss_customer_sk#2, null, null, null, 1], [null, ss_item_sk#1, null, null, 2], [null, null, ss_store_sk#6, null, 3], [null, null, null, ss_promo_sk#7, 4]], [ss_customer_sk#60, ss_item_sk#61, ss_store_sk#62, ss_promo_sk#63, gid#59]";
+
+        const result = parseExpand(input);
+
+        expect(result).toEqual({
+            fields: [
+                "ss_customer_sk",
+                "ss_item_sk",
+                "ss_store_sk",
+                "ss_promo_sk"
+            ],
+            idField: "gid"
+        });
+    });
+
+    test("should handle expand plan with different field names", () => {
+        const input =
+            "Expand [[field1#1, null, 0], [null, field2#2, 1]], [output1#10, output2#11, group_id#12]";
+
+        const result = parseExpand(input);
+
+        expect(result).toEqual({
+            fields: [
+                "output1",
+                "output2"
+            ],
+            idField: "group_id"
+        });
+    });
+
+    test("should throw error for invalid format", () => {
+        const input = "Invalid expand format";
+
+        expect(() => parseExpand(input)).toThrow("Invalid Expand plan format: could not find output fields");
+    });
+});

spark-ui/src/reducers/PlanParsers/ExpandParser.ts

@@ -0,0 +1,38 @@
+import { ParsedExpandPlan } from "../../interfaces/AppStore";
+import { bracedSplit, hashNumbersRemover } from "./PlanParserUtils";
+
+export function parseExpand(input: string): ParsedExpandPlan {
+    // Clean input by removing hash numbers
+    const cleanInput = hashNumbersRemover(input);
+
+    // Parse the Expand plan: "Expand [[ss_customer_sk#2, null, null, null, 1], [null, ss_item_sk#1, null, null, 2], [null, null, ss_store_sk#6, null, 3], [null, null, null, ss_promo_sk#7, 4]], [ss_customer_sk#60, ss_item_sk#61, ss_store_sk#62, ss_promo_sk#63, gid#59]"
+    // We need to extract the output fields from the final bracketed section after the nested arrays
+
+    // Look for the pattern: ]], [fields...]
+    const finalBracketRegex = /\]\],\s*\[([^\]]+)\]$/;
+    const match = cleanInput.match(finalBracketRegex);
+
+    if (!match) {
+        throw new Error("Invalid Expand plan format: could not find output fields");
+    }
+
+    const fieldsStr = match[1];
+
+    // Split the fields by comma and clean them up
+    const allFields = bracedSplit(fieldsStr)
+        .map(field => field.trim())
+        .filter(field => field.length > 0);
+
+    if (allFields.length === 0) {
+        throw new Error("Invalid Expand plan format: no output fields found");
+    }
+
+    // Extract the last field as idField and remove it from the fields array
+    const idField = allFields[allFields.length - 1];
+    const fields = allFields.slice(0, -1);
+
+    return {
+        fields: fields,
+        idField: idField
+    };
+}

spark-ui/src/reducers/SqlReducer.ts

@@ -26,6 +26,7 @@ import { findLastNodeWithInputRows, generateGraph, getRowsFromMetrics } from "./
 import { parseCoalesce } from "./PlanParsers/CoalesceParser";
 import { parseCollectLimit } from "./PlanParsers/CollectLimitParser";
 import { parseExchange } from "./PlanParsers/ExchangeParser";
+import { parseExpand } from "./PlanParsers/ExpandParser";
 import { parseFilter } from "./PlanParsers/FilterParser";
 import { parseGenerate } from "./PlanParsers/GenerateParser";
 import { parseJoin } from "./PlanParsers/JoinParser";
@@ -168,6 +169,11 @@ export function parseNodePlan(
           type: "Generate",
           plan: parseGenerate(plan.planDescription),
         };
+      case "Expand":
+        return {
+          type: "Expand",
+          plan: parseExpand(plan.planDescription),
+        };
     }
     if (node.nodeName.includes("Scan")) {
       return {
@@ -733,6 +739,36 @@ function addGenerateMetrics(
   return null;
 }
 
+function addExpandMetrics(
+  node: EnrichedSqlNode,
+  updatedMetrics: EnrichedSqlMetric[],
+  graph: Graph,
+  allNodes: EnrichedSqlNode[],
+): EnrichedSqlMetric | null {
+  if (node.nodeName === "Expand") {
+    const inputNode = findLastNodeWithInputRows(node, graph, allNodes);
+    if (!inputNode) {
+      return null;
+    }
+
+    const inputRows = getRowsFromMetrics(inputNode.metrics);
+    if (inputRows === null || inputRows === 0) {
+      return null;
+    }
+
+    const outputRows = getRowsFromMetrics(updatedMetrics);
+    if (outputRows === null) {
+      return null;
+    }
+
+    const ratio = outputRows / inputRows;
+    const ratioFormatted = ratio.toFixed(2);
+
+    return { name: `Expand Ratio`, value: `${ratioFormatted}X` };
+  }
+  return null;
+}
+
 function updateNodeMetrics(
   node: EnrichedSqlNode,
   metrics: EnrichedSqlMetric[],
@@ -744,6 +780,7 @@ function updateNodeMetrics(
   const crossJoinFilterRatio = addCrossJoinFilterRatioMetric(node, updatedOriginalMetrics, graph, allNodes);
   const joinMetrics = addJoinMetrics(node, updatedOriginalMetrics, graph, allNodes);
   const generateMetrics = addGenerateMetrics(node, updatedOriginalMetrics, graph, allNodes);
+  const expandMetrics = addExpandMetrics(node, updatedOriginalMetrics, graph, allNodes);
   return [
     ...updatedOriginalMetrics,
     ...(filterRatio !== null
@@ -758,6 +795,9 @@ function updateNodeMetrics(
     ...(generateMetrics !== null
       ? [generateMetrics]
       : []),
+    ...(expandMetrics !== null
+      ? [expandMetrics]
+      : []),
   ];
 }
 

spark-ui/src/reducers/SqlReducerUtils.ts

@@ -157,6 +157,7 @@ const nodeTypeDict: Record<string, NodeType> = {
   FlatMapGroupsInPandas: "transformation",
   BatchEvalPython: "transformation",
   Generate: "transformation",
+  Expand: "transformation",
 };
 
 const nodeRenamerDict: Record<string, string> = {
@@ -225,6 +226,7 @@ const nodeRenamerDict: Record<string, string> = {
   ArrowEvalPython: "Select (with Arrow)",
   FlatMapGroupsInPandas: "Select Flat (with Pandas)",
   BatchEvalPython: "Run Python UDF",
+  Expand: "Expand",
 };
 
 export function extractTotalFromStatisticsMetric(
@@ -280,6 +282,8 @@ export function nodeEnrichedNameBuilder(
           return plan.plan.operation;
         }
         return "Generate";
+      case "Expand":
+        return "Expand";
       case "Exchange":
         if (plan.plan.isBroadcast) {
           return "Broadcast";