From b05f3423f802cb27ac7f7bf8d563c3040a5e1221 Mon Sep 17 00:00:00 2001
From: Michael Johns <mjohns@databricks.com>
Date: Wed, 3 Jan 2024 14:36:26 -0500
Subject: [PATCH] PathUtils handle dbfs:/Volumes paths.

---
 python/mosaic/config/config.py                |  1 -
 .../labs/mosaic/utils/PathUtils.scala         | 26 +++++++++++++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/python/mosaic/config/config.py b/python/mosaic/config/config.py
index bc5f80c9f..a59979be6 100644
--- a/python/mosaic/config/config.py
+++ b/python/mosaic/config/config.py
@@ -10,4 +10,3 @@
 display_handler: DisplayHandler
 ipython_hook: InteractiveShell
 notebook_utils = None
-default_gdal_init_script_path: str = "/dbfs/FileStore/geospatial/mosaic/gdal/"
diff --git a/src/main/scala/com/databricks/labs/mosaic/utils/PathUtils.scala b/src/main/scala/com/databricks/labs/mosaic/utils/PathUtils.scala
index f3fb9d7b9..965e3fdde 100644
--- a/src/main/scala/com/databricks/labs/mosaic/utils/PathUtils.scala
+++ b/src/main/scala/com/databricks/labs/mosaic/utils/PathUtils.scala
@@ -9,7 +9,18 @@ import java.nio.file.{Files, Paths}
 object PathUtils {
 
     def getCleanPath(path: String): String = {
-        val cleanPath = path.replace("file:/", "/").replace("dbfs:/", "/dbfs/")
+        //val cleanPath = path.replace("file:/", "/").replace("dbfs:/", "/dbfs/")
+        val cleanPath = {
+            if (path.startsWith("file:/")) {
+                path.replace("file:/", "/")
+            } else if (path.startsWith("dbfs:/Volumes")) {
+                path.replace("dbfs:/Volumes", "/Volumes")
+            } else if (path.startsWith("dbfs:/")) {
+                path.replace("dbfs:/", "/dbfs/")
+            } else {
+                path
+            }
+        }
         if (cleanPath.endsWith(".zip") || cleanPath.contains(".zip:")) {
             getZipPath(cleanPath)
         } else {
@@ -59,7 +70,18 @@ object PathUtils {
 
     def copyToTmp(inPath: String): String = {
         val cleanPath = getCleanPath(inPath)
-        val copyFromPath = inPath.replace("file:/", "/").replace("dbfs:/", "/dbfs/")
+        //val copyFromPath = inPath.replace("file:/", "/").replace("dbfs:/", "/dbfs/")
+        val copyFromPath = {
+            if (inPath.startsWith("file:/")) {
+                inPath.replace("file:/", "/")
+            } else if (inPath.startsWith("dbfs:/Volumes")) {
+                inPath.replace("dbfs:/Volumes", "/Volumes")
+            } else if (inPath.startsWith("dbfs:/")) {
+                inPath.replace("dbfs:/", "/dbfs/")
+            } else {
+                inPath
+            }
+        }
         val driver = MosaicRasterGDAL.identifyDriver(cleanPath)
         val extension = if (inPath.endsWith(".zip")) "zip" else GDAL.getExtension(driver)
         val tmpPath = createTmpFilePath(extension)