From a0a5ca82dda44ca309d0c62abe7eac8b4fff676e Mon Sep 17 00:00:00 2001 From: "Mateusz \"Serafin\" Gajewski" Date: Thu, 28 Nov 2024 16:28:22 +0100 Subject: [PATCH 1/2] 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 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 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 List> column(T... values) { return Arrays.stream(values) - .map(List::of) - .collect(toImmutableList()); + // Allow nulls in values + .map(value -> { + List list = new ArrayList<>(); + list.add(value); + return list; + }) + .collect(toList()); } private static List array(T... values) From 542cb2c3b0199c1fb341adb9503329a53589b216 Mon Sep 17 00:00:00 2001 From: "Mateusz \"Serafin\" Gajewski" Date: Thu, 28 Nov 2024 16:30:13 +0100 Subject: [PATCH 2/2] Rename test class --- ...nQueryDataEncoding.java => TestJsonEncodingUtils.java} | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) rename core/trino-main/src/test/java/io/trino/server/protocol/{spooling/TestJsonQueryDataEncoding.java => TestJsonEncodingUtils.java} (99%) 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/TestJsonEncodingUtils.java similarity index 99% rename from core/trino-main/src/test/java/io/trino/server/protocol/spooling/TestJsonQueryDataEncoding.java rename to core/trino-main/src/test/java/io/trino/server/protocol/TestJsonEncodingUtils.java index 447ced92d6c1f..094e77f3d4dc9 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/TestJsonEncodingUtils.java @@ -11,7 +11,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package io.trino.server.protocol.spooling; +package io.trino.server.protocol; import com.google.common.collect.ImmutableList; import io.trino.client.Column; @@ -19,7 +19,7 @@ import io.trino.client.Row; import io.trino.client.spooling.DataAttributes; import io.trino.client.spooling.encoding.JsonQueryDataDecoder; -import io.trino.server.protocol.OutputColumn; +import io.trino.server.protocol.spooling.QueryDataEncoder; import io.trino.server.protocol.spooling.encoding.JsonQueryDataEncoder; import io.trino.spi.Page; import io.trino.spi.block.ArrayBlockBuilder; @@ -56,7 +56,7 @@ import static io.trino.block.BlockAssertions.createTinyintsBlock; import static io.trino.block.BlockAssertions.createTypedLongsBlock; import static io.trino.server.protocol.ProtocolUtil.createColumn; -import static io.trino.server.protocol.spooling.TestJsonQueryDataEncoding.TypedColumn.typed; +import static io.trino.server.protocol.TestJsonEncodingUtils.TypedColumn.typed; import static io.trino.spi.type.BigintType.BIGINT; import static io.trino.spi.type.BooleanType.BOOLEAN; import static io.trino.spi.type.DoubleType.DOUBLE; @@ -72,7 +72,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; -public class TestJsonQueryDataEncoding +public class TestJsonEncodingUtils { protected QueryDataDecoder createDecoder(List columns) {