From bc9a630f664296ecb3e5cc563d5d81808fa91744 Mon Sep 17 00:00:00 2001
From: "Mateusz \"Serafin\" Gajewski" <github@wendigo.pl>
Date: Thu, 28 Nov 2024 16:28:22 +0100
Subject: [PATCH] Write varchar slice directly to JSON

---
 .../server/protocol/JsonEncodingUtils.java    |  2 +-
 .../spooling/TestJsonQueryDataEncoding.java   | 35 +++++++++++++++----
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/core/trino-main/src/main/java/io/trino/server/protocol/JsonEncodingUtils.java b/core/trino-main/src/main/java/io/trino/server/protocol/JsonEncodingUtils.java
index 30abaf5407a01..d930cbb82cd31 100644
--- a/core/trino-main/src/main/java/io/trino/server/protocol/JsonEncodingUtils.java
+++ b/core/trino-main/src/main/java/io/trino/server/protocol/JsonEncodingUtils.java
@@ -264,7 +264,7 @@ public void encode(JsonGenerator generator, ConnectorSession session, Block bloc
                 return;
             }
             Slice slice = VARCHAR.getSlice(block, position);
-            generator.writeString(slice.toStringUtf8());
+            generator.writeUTF8String(slice.byteArray(), slice.byteArrayOffset(), slice.length());
         }
     }
 
diff --git a/core/trino-main/src/test/java/io/trino/server/protocol/spooling/TestJsonQueryDataEncoding.java b/core/trino-main/src/test/java/io/trino/server/protocol/spooling/TestJsonQueryDataEncoding.java
index fbdac9f95e080..447ced92d6c1f 100644
--- a/core/trino-main/src/test/java/io/trino/server/protocol/spooling/TestJsonQueryDataEncoding.java
+++ b/core/trino-main/src/test/java/io/trino/server/protocol/spooling/TestJsonQueryDataEncoding.java
@@ -37,12 +37,12 @@
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.airlift.slice.Slices.utf8Slice;
 import static io.trino.SessionTestUtils.TEST_SESSION;
@@ -68,6 +68,7 @@
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
+import static java.util.stream.Collectors.toList;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
@@ -291,10 +292,27 @@ public void testVarcharSerialization()
             throws IOException
     {
         List<TypedColumn> columns = ImmutableList.of(typed("col0", VARCHAR));
-        Page page = page(createStringsBlock("ala", "ma", "kota", "a", "kot", "ma", "ale", ""));
+        Page page = page(createStringsBlock("ala", "ma", "kota", "a", "kot", "ma", "ale", "", null));
 
-        assertThat(roundTrip(columns, page, "[[\"ala\"],[\"ma\"],[\"kota\"],[\"a\"],[\"kot\"],[\"ma\"],[\"ale\"],[\"\"]]"))
-                .isEqualTo(column("ala", "ma", "kota", "a", "kot", "ma", "ale", ""));
+        assertThat(roundTrip(columns, page, "[[\"ala\"],[\"ma\"],[\"kota\"],[\"a\"],[\"kot\"],[\"ma\"],[\"ale\"],[\"\"],[null]]"))
+                .isEqualTo(column("ala", "ma", "kota", "a", "kot", "ma", "ale", "", null));
+    }
+
+    @Test
+    public void testVarcharUtf8Serialization()
+            throws IOException
+    {
+        List<TypedColumn> columns = ImmutableList.of(typed("col0", VARCHAR));
+        Page page = page(createStringsBlock(
+                "数据应用",
+                "\"quoted\"",
+                "zażółć gęślą jaźń",
+                "\0\0\0",  // garbage in, garbage out
+                "\r\t\n",
+                "\uD83E\uDD83"));
+
+        assertThat(roundTrip(columns, page, "[[\"数据应用\"],[\"\\\"quoted\\\"\"],[\"zażółć gęślą jaźń\"],[\"\\u0000\\u0000\\u0000\"],[\"\\r\\t\\n\"],[\"🦃\"]]"))
+                .isEqualTo(column("数据应用", "\"quoted\"", "zażółć gęślą jaźń", "\0\0\0", "\r\t\n", "🦃"));
     }
 
     @Test
@@ -564,8 +582,13 @@ private static Page page(Block... blocks)
     private static <T> List<List<T>> column(T... values)
     {
         return Arrays.stream(values)
-                .map(List::of)
-                .collect(toImmutableList());
+                // Allow nulls in values
+                .map(value -> {
+                    List<T> list = new ArrayList<>();
+                    list.add(value);
+                    return list;
+                })
+                .collect(toList());
     }
 
     private static <T> List<T> array(T... values)