[SPARK-52320] Add ColumnNotFound/InvalidViewName/TableOrViewAlreadyExists to SparkConnectError

Author: dongjoon-hyun

Sources/SparkConnect/DataFrame.swift

@@ -309,6 +309,8 @@ public actor DataFrame: Sendable {
           throw SparkConnectError.SchemaNotFound
         case let m where m.contains("TABLE_OR_VIEW_NOT_FOUND"):
           throw SparkConnectError.TableOrViewNotFound
+        case let m where m.contains("UNRESOLVED_COLUMN.WITH_SUGGESTION"):
+          throw SparkConnectError.ColumnNotFound
         default:
           throw error
         }

Sources/SparkConnect/SparkConnectClient.swift

@@ -127,7 +127,20 @@ public actor SparkConnectClient {
       ),
       interceptors: self.intercepters
     ) { client in
-      return try await f(client)
+      do {
+        return try await f(client)
+      } catch let error as RPCError where error.code == .internalError {
+        switch error.message {
+        case let m where m.contains("TABLE_OR_VIEW_ALREADY_EXISTS"):
+          throw SparkConnectError.TableOrViewAlreadyExists
+        case let m where m.contains("TABLE_OR_VIEW_NOT_FOUND"):
+          throw SparkConnectError.TableOrViewNotFound
+        case let m where m.contains("Invalid view name:"):
+          throw SparkConnectError.InvalidViewName
+        default:
+          throw error
+        }
+      }
     }
   }
 

Sources/SparkConnect/SparkConnectError.swift

@@ -20,10 +20,13 @@
 /// A enum for ``SparkConnect`` package errors
 public enum SparkConnectError: Error {
   case CatalogNotFound
+  case ColumnNotFound
   case InvalidArgument
   case InvalidSessionID
   case InvalidType
+  case InvalidViewName
   case SchemaNotFound
+  case TableOrViewAlreadyExists
   case TableOrViewNotFound
   case UnsupportedOperation
 }

Tests/SparkConnectTests/CatalogTests.swift

@@ -205,12 +205,12 @@ struct CatalogTests {
       try await spark.range(1).createTempView(viewName)
       #expect(try await spark.catalog.tableExists(viewName))
 
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewAlreadyExists) {
         try await spark.range(1).createTempView(viewName)
       }
     })
 
-    try await #require(throws: Error.self) {
+    try await #require(throws: SparkConnectError.InvalidViewName) {
       try await spark.range(1).createTempView("invalid view name")
     }
 
@@ -228,7 +228,7 @@ struct CatalogTests {
       try await spark.range(1).createOrReplaceTempView(viewName)
     })
 
-    try await #require(throws: Error.self) {
+    try await #require(throws: SparkConnectError.InvalidViewName) {
       try await spark.range(1).createOrReplaceTempView("invalid view name")
     }
 
@@ -244,13 +244,13 @@ struct CatalogTests {
       try await spark.range(1).createGlobalTempView(viewName)
       #expect(try await spark.catalog.tableExists("global_temp.\(viewName)"))
 
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewAlreadyExists) {
         try await spark.range(1).createGlobalTempView(viewName)
       }
     })
     #expect(try await spark.catalog.tableExists("global_temp.\(viewName)") == false)
 
-    try await #require(throws: Error.self) {
+    try await #require(throws: SparkConnectError.InvalidViewName) {
       try await spark.range(1).createGlobalTempView("invalid view name")
     }
 
@@ -269,7 +269,7 @@ struct CatalogTests {
     })
     #expect(try await spark.catalog.tableExists("global_temp.\(viewName)") == false)
 
-    try await #require(throws: Error.self) {
+    try await #require(throws: SparkConnectError.InvalidViewName) {
       try await spark.range(1).createOrReplaceGlobalTempView("invalid view name")
     }
 

Tests/SparkConnectTests/DataFrameTests.swift

@@ -235,10 +235,10 @@ struct DataFrameTests {
   @Test
   func selectInvalidColumn() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
-    try await #require(throws: Error.self) {
+    try await #require(throws: SparkConnectError.ColumnNotFound) {
       try await spark.range(1).select("invalid").schema
     }
-    try await #require(throws: Error.self) {
+    try await #require(throws: SparkConnectError.ColumnNotFound) {
       try await spark.range(1).select("id + 1").schema
     }
     await spark.stop()

Tests/SparkConnectTests/DataFrameWriterTests.swift

@@ -112,7 +112,7 @@ struct DataFrameWriterTests {
       try await spark.range(1).write.saveAsTable(tableName)
       #expect(try await spark.read.table(tableName).count() == 1)
 
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewAlreadyExists) {
         try await spark.range(1).write.saveAsTable(tableName)
       }
 
@@ -130,8 +130,7 @@ struct DataFrameWriterTests {
     let spark = try await SparkSession.builder.getOrCreate()
     let tableName = "TABLE_" + UUID().uuidString.replacingOccurrences(of: "-", with: "")
     try await SQLHelper.withTable(spark, tableName)({
-      // Table doesn't exist.
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewNotFound) {
         try await spark.range(1).write.insertInto(tableName)
       }
 

Tests/SparkConnectTests/DataFrameWriterV2Tests.swift

@@ -34,7 +34,7 @@ struct DataFrameWriterV2Tests {
       let write = try await spark.range(2).writeTo(tableName).using("orc")
       try await write.create()
       #expect(try await spark.table(tableName).count() == 2)
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewAlreadyExists) {
         try await write.create()
       }
     })

Tests/SparkConnectTests/MergeIntoWriterTests.swift

@@ -31,10 +31,10 @@ struct MergeIntoWriterTests {
     let tableName = "TABLE_" + UUID().uuidString.replacingOccurrences(of: "-", with: "")
     try await SQLHelper.withTable(spark, tableName)({
       let mergeInto = try await spark.range(1).mergeInto(tableName, "true")
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewNotFound) {
         try await mergeInto.whenMatched().delete().merge()
       }
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewNotFound) {
         try await mergeInto.whenMatched("true").delete().merge()
       }
     })
@@ -47,10 +47,10 @@ struct MergeIntoWriterTests {
     let tableName = "TABLE_" + UUID().uuidString.replacingOccurrences(of: "-", with: "")
     try await SQLHelper.withTable(spark, tableName)({
       let mergeInto = try await spark.range(1).mergeInto(tableName, "true")
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewNotFound) {
         try await mergeInto.whenNotMatched().insertAll().merge()
       }
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewNotFound) {
         try await mergeInto.whenNotMatched("true").insertAll().merge()
       }
     })
@@ -63,10 +63,10 @@ struct MergeIntoWriterTests {
     let tableName = "TABLE_" + UUID().uuidString.replacingOccurrences(of: "-", with: "")
     try await SQLHelper.withTable(spark, tableName)({
       let mergeInto = try await spark.range(1).mergeInto(tableName, "true")
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewNotFound) {
         try await mergeInto.whenNotMatchedBySource().delete().merge()
       }
-      try await #require(throws: Error.self) {
+      try await #require(throws: SparkConnectError.TableOrViewNotFound) {
         try await mergeInto.whenNotMatchedBySource("true").delete().merge()
       }
     })