From 5cc8dcb6afe55076d77e7c57574b57045ace163c Mon Sep 17 00:00:00 2001
From: Robin Tang <rtang.cs@gmail.com>
Date: Tue, 17 Dec 2024 21:09:53 -0800
Subject: [PATCH] [Databricks] Preprocess Arrays before `COPY INTO` (#1084)

---
 clients/databricks/store.go      | 32 +++++++++++++++++++++++++++++++-
 lib/typing/values/string.go      |  8 ++++++++
 lib/typing/values/string_test.go | 15 ++++++++++++---
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/clients/databricks/store.go b/clients/databricks/store.go
index 02df68b94..13d2b1618 100644
--- a/clients/databricks/store.go
+++ b/clients/databricks/store.go
@@ -7,6 +7,7 @@ import (
 	"log/slog"
 	"os"
 	"path/filepath"
+	"strings"
 
 	_ "github.com/databricks/databricks-sql-go"
 	"github.com/databricks/databricks-sql-go/driverctx"
@@ -120,9 +121,38 @@ func (s Store) PrepareTemporaryTable(ctx context.Context, tableData *optimizatio
 		}
 	}()
 
+	var ordinalColumns []string
+	for idx, column := range tableData.ReadOnlyInMemoryCols().ValidColumns() {
+		ordinalColumn := fmt.Sprintf("_c%d", idx)
+		switch column.KindDetails.Kind {
+		case typing.Array.Kind:
+			ordinalColumn = fmt.Sprintf(`PARSE_JSON(%s)`, ordinalColumn)
+		}
+
+		ordinalColumns = append(ordinalColumns, ordinalColumn)
+	}
+
 	// Copy file from DBFS -> table via COPY INTO, ref: https://docs.databricks.com/en/sql/language-manual/delta-copy-into.html
 	// We'll need \\\\N here because we need to string escape.
-	copyCommand := fmt.Sprintf(`COPY INTO %s BY POSITION FROM '%s' FILEFORMAT = CSV FORMAT_OPTIONS ('escape' = '"', 'delimiter' = '\t', 'header' = 'false', 'nullValue' = '\\\\N')`, tempTableID.FullyQualifiedName(), file.DBFSFilePath())
+	copyCommand := fmt.Sprintf(`
+COPY INTO %s
+BY POSITION
+FROM (
+    SELECT %s FROM '%s'
+)
+FILEFORMAT = CSV
+FORMAT_OPTIONS (
+    'escape' = '"', 
+    'delimiter' = '\t', 
+    'header' = 'false', 
+    'nullValue' = '\\\\N'
+);`,
+		// COPY INTO
+		tempTableID.FullyQualifiedName(),
+		// SELECT columns FROM file
+		strings.Join(ordinalColumns, ", "), file.DBFSFilePath(),
+	)
+
 	if _, err = s.ExecContext(ctx, copyCommand); err != nil {
 		return fmt.Errorf("failed to run COPY INTO for temporary table: %w", err)
 	}
diff --git a/lib/typing/values/string.go b/lib/typing/values/string.go
index 544be3a9e..8114ea610 100644
--- a/lib/typing/values/string.go
+++ b/lib/typing/values/string.go
@@ -97,6 +97,14 @@ func ToString(colVal any, colKind typing.KindDetails) (string, error) {
 			}
 		}
 	case typing.Array.Kind:
+		// If the column value is TOASTED, we should return an array with the TOASTED placeholder
+		// We're doing this to make sure that the value matches the schema.
+		if stringValue, ok := colVal.(string); ok {
+			if stringValue == constants.ToastUnavailableValuePlaceholder {
+				return fmt.Sprintf(`["%s"]`, constants.ToastUnavailableValuePlaceholder), nil
+			}
+		}
+
 		colValBytes, err := json.Marshal(colVal)
 		if err != nil {
 			return "", err
diff --git a/lib/typing/values/string_test.go b/lib/typing/values/string_test.go
index 8db4751c4..5ad48a2af 100644
--- a/lib/typing/values/string_test.go
+++ b/lib/typing/values/string_test.go
@@ -172,9 +172,18 @@ func TestToString(t *testing.T) {
 	}
 	{
 		// Array
-		val, err := ToString([]string{"foo", "bar"}, typing.Array)
-		assert.NoError(t, err)
-		assert.Equal(t, `["foo","bar"]`, val)
+		{
+			// Normal arrays
+			val, err := ToString([]string{"foo", "bar"}, typing.Array)
+			assert.NoError(t, err)
+			assert.Equal(t, `["foo","bar"]`, val)
+		}
+		{
+			// Toasted array
+			val, err := ToString(constants.ToastUnavailableValuePlaceholder, typing.Array)
+			assert.NoError(t, err)
+			assert.Equal(t, `["__debezium_unavailable_value"]`, val)
+		}
 	}
 	{
 		// Integer column