optimize UTF-8 fallback after ASCII prefix

Author: jerrinot

core/src/main/java/io/questdb/client/cutlass/http/client/WebSocketSendBuffer.java

@@ -334,7 +334,7 @@ public void putUtf8(String value) {
         int charLen = value.length();
         ensureCapacity(charLen);
 
-        // Single-pass ASCII path. Non-ASCII falls back to the shared UTF-8 utility.
+        // Single-pass for ASCII. Mixed strings keep the ASCII prefix and resume UTF-8 encoding.
         long addr = bufPtr + writePos;
         int i = 0;
         for (; i < charLen; i++) {
@@ -348,9 +348,9 @@ public void putUtf8(String value) {
         if (i == charLen) {
             writePos += charLen;
         } else {
-            int utf8Len = NativeBufferWriter.utf8Length(value);
-            ensureCapacity(utf8Len);
-            writePos += Utf8s.strCpyUtf8(value, bufPtr + writePos, utf8Len);
+            int utf8Len = Utf8s.utf8Bytes(value, i, charLen);
+            ensureCapacity(i + utf8Len);
+            writePos += i + Utf8s.strCpyUtf8(value, i, bufPtr + writePos + i, utf8Len);
         }
     }
 

core/src/main/java/io/questdb/client/cutlass/qwp/client/NativeBufferWriter.java

@@ -289,10 +289,9 @@ public void putUtf8(String value) {
             // All ASCII — done in a single pass
             position += charLen;
         } else {
-            // Non-ASCII — fall back to two-pass (re-encodes from start)
-            int utf8Len = utf8Length(value);
-            ensureCapacity(utf8Len);
-            encodeUtf8(value, utf8Len);
+            int utf8Len = Utf8s.utf8Bytes(value, i, charLen);
+            ensureCapacity(i + utf8Len);
+            position += i + Utf8s.strCpyUtf8(value, i, bufferPtr + position + i, utf8Len);
         }
     }
 

core/src/main/java/io/questdb/client/std/str/Utf8s.java

@@ -181,8 +181,12 @@ public static void strCpyAscii(@NotNull CharSequence asciiSrc, int srcLo, int sr
      * @return the number of UTF-8 bytes written
      */
     public static int strCpyUtf8(@NotNull CharSequence src, long destAddr, int maxBytes) {
+        return strCpyUtf8(src, 0, destAddr, maxBytes);
+    }
+
+    public static int strCpyUtf8(@NotNull CharSequence src, int srcLo, long destAddr, int maxBytes) {
         int pos = 0;
-        for (int i = 0, n = src.length(); i < n; i++) {
+        for (int i = srcLo, n = src.length(); i < n; i++) {
             char c = src.charAt(i);
             if (c < 0x80) {
                 if (pos + 1 > maxBytes) {
@@ -263,18 +267,20 @@ public static String stringFromUtf8Bytes(@NotNull Utf8Sequence seq) {
     }
 
     public static int utf8Bytes(@NotNull CharSequence sequence) {
-        int count = 0;
-        int len = sequence.length();
+        return utf8Bytes(sequence, 0, sequence.length());
+    }
 
-        for (int i = 0; i < len; i++) {
+    public static int utf8Bytes(@NotNull CharSequence sequence, int lo, int hi) {
+        int count = 0;
+        for (int i = lo; i < hi; i++) {
             char ch = sequence.charAt(i);
             if (ch < 0x80) {
                 count++;
             } else if (ch < 0x800) {
                 count += 2;
             } else if (Character.isSurrogate(ch)) {
                 if (Character.isHighSurrogate(ch)) {
-                    if (i + 1 < len && Character.isLowSurrogate(sequence.charAt(i + 1))) {
+                    if (i + 1 < hi && Character.isLowSurrogate(sequence.charAt(i + 1))) {
                         count += 4;
                         i++;
                     } else {

core/src/test/java/io/questdb/client/test/cutlass/http/client/WebSocketSendBufferTest.java

@@ -26,9 +26,10 @@
 
 import io.questdb.client.cutlass.http.client.WebSocketSendBuffer;
 import io.questdb.client.std.Unsafe;
-import static io.questdb.client.test.tools.TestUtils.assertMemoryLeak;
+import io.questdb.client.std.str.Utf8s;
 import org.junit.Test;
 
+import static io.questdb.client.test.tools.TestUtils.assertMemoryLeak;
 import static org.junit.Assert.assertEquals;
 
 public class WebSocketSendBufferTest {
@@ -46,4 +47,19 @@ public void testPutUtf8InvalidSurrogatePair() throws Exception {
             }
         });
     }
+
+    @Test
+    public void testPutUtf8MixedAsciiAndNonAsciiAfterGrow() throws Exception {
+        assertMemoryLeak(() -> {
+            try (WebSocketSendBuffer buf = new WebSocketSendBuffer(8)) {
+                String value = "abcdefghijklmnop世界世界世界世界世界世界世界世界世界世界";
+
+                buf.putUtf8(value);
+
+                int utf8Len = Utf8s.utf8Bytes(value);
+                assertEquals(utf8Len, buf.getWritePos());
+                assertEquals(value, Utf8s.stringFromUtf8Bytes(buf.getBufferPtr(), buf.getBufferPtr() + utf8Len));
+            }
+        });
+    }
 }

core/src/test/java/io/questdb/client/test/cutlass/qwp/client/NativeBufferWriterTest.java

@@ -26,6 +26,7 @@
 
 import io.questdb.client.cutlass.qwp.client.NativeBufferWriter;
 import io.questdb.client.std.Unsafe;
+import io.questdb.client.std.str.Utf8s;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -522,6 +523,21 @@ public void testWriteUtf8MixedAsciiAndNonAscii() throws Exception {
         });
     }
 
+    @Test
+    public void testWriteUtf8MixedAsciiAndNonAsciiAfterGrow() throws Exception {
+        assertMemoryLeak(() -> {
+            try (NativeBufferWriter writer = new NativeBufferWriter(8)) {
+                String value = "abcdefghijklmnop世界世界世界世界世界世界世界世界世界世界";
+
+                writer.putUtf8(value);
+
+                int utf8Len = Utf8s.utf8Bytes(value);
+                Assert.assertEquals(utf8Len, writer.getPosition());
+                Assert.assertEquals(value, Utf8s.stringFromUtf8Bytes(writer.getBufferPtr(), writer.getBufferPtr() + utf8Len));
+            }
+        });
+    }
+
     @Test
     public void testWriteUtf8Ascii() throws Exception {
         assertMemoryLeak(() -> {

core/src/test/java/io/questdb/client/test/std/str/Utf8sTest.java

@@ -270,6 +270,27 @@ public void testStrCpyUtf8() {
         }
     }
 
+    @Test
+    public void testStrCpyUtf8FromOffset() {
+        final int bufSize = 64;
+        long mem = Unsafe.malloc(bufSize, MemoryTag.NATIVE_DEFAULT);
+        try {
+            final byte sentinel = 0x5A;
+            String value = "ascii-prefixé世\uD83D\uDE00\uD800x";
+            int lo = "ascii-prefix".length();
+            int expectedBytes = Utf8s.utf8Bytes(value, lo, value.length());
+
+            fill(mem, bufSize, sentinel);
+            Assert.assertEquals(expectedBytes, Utf8s.strCpyUtf8(value, lo, mem, expectedBytes));
+            Assert.assertEquals("é世\uD83D\uDE00?x", readUtf8(mem, expectedBytes));
+            for (int i = expectedBytes; i < bufSize; i++) {
+                Assert.assertEquals("write past copied bytes at offset " + i, sentinel, Unsafe.getUnsafe().getByte(mem + i));
+            }
+        } finally {
+            Unsafe.free(mem, bufSize, MemoryTag.NATIVE_DEFAULT);
+        }
+    }
+
     @Test
     public void testStrCpyUtf8DoesNotSplitCharactersAtLimit() {
         final int bufSize = 16;
@@ -319,6 +340,14 @@ public void testUtf8Bytes() {
         Assert.assertEquals(1, Utf8s.utf8Bytes("\uDE00"));
     }
 
+    @Test
+    public void testUtf8BytesRange() {
+        String value = "xxAé世\uD83D\uDE00yy";
+        Assert.assertEquals(10, Utf8s.utf8Bytes(value, 2, 7));
+        Assert.assertEquals(1, Utf8s.utf8Bytes(value, 5, 6));
+        Assert.assertEquals(4, Utf8s.utf8Bytes(value, 5, 7));
+    }
+
     @Test
     public void testUtf8BytesWithLimit() {
         Assert.assertEquals(0, Utf8s.utf8Bytes("hello", 0));