xiangfu0 commented on code in PR #18953: URL: https://github.com/apache/pinot/pull/18953#discussion_r3556309771
########## pinot-plugins/pinot-input-format/pinot-json/src/main/java/org/apache/pinot/plugin/inputformat/json/format/SqliteJsonbPayloadParser.java: ########## @@ -0,0 +1,383 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.plugin.inputformat.json.format; + +import com.google.common.collect.Maps; +import java.math.BigInteger; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + + +/// Parses the <a href="https://sqlite.org/jsonb.html";>SQLite JSONB</a> binary format (SQLite 3.45+). +/// +/// Every element is a 1–9 byte header followed by a payload. The header's first byte packs the element type in +/// the low nibble and a payload-size descriptor in the high nibble; descriptors 12–15 mean the size is a big- +/// endian `uint8` / `uint16` / `uint32` / `uint64` in the following bytes, and 0–11 are the size itself. +/// +/// Numbers are stored as their ASCII text, so integers narrow to `Integer` / `Long` / `BigInteger` and floats +/// to `Double`, matching the text-JSON value contract. `TEXTJ` strings/labels are JSON-unescaped (invalid +/// escapes are rejected); `TEXT5` additionally accepts JSON5 escapes. +/// +/// The decoder expects one row per payload, so the top-level element must be an `OBJECT`. +class SqliteJsonbPayloadParser implements JsonPayloadParser { + + // Element types (low nibble of the header's first byte). + private static final int TYPE_NULL = 0; + private static final int TYPE_TRUE = 1; + private static final int TYPE_FALSE = 2; + private static final int TYPE_INT = 3; // canonical decimal integer, ASCII + private static final int TYPE_INT5 = 4; // JSON5 integer (hex / leading sign), ASCII + private static final int TYPE_FLOAT = 5; // canonical float, ASCII + private static final int TYPE_FLOAT5 = 6; // JSON5 float, ASCII + private static final int TYPE_TEXT = 7; // raw text, no escapes + private static final int TYPE_TEXTJ = 8; // text with JSON escapes + private static final int TYPE_TEXT5 = 9; // text with JSON5 escapes + private static final int TYPE_TEXTRAW = 10; // raw text needing quoting, no escapes + private static final int TYPE_ARRAY = 11; + private static final int TYPE_OBJECT = 12; + + private static final BigInteger INT_MIN = BigInteger.valueOf(Integer.MIN_VALUE); + private static final BigInteger INT_MAX = BigInteger.valueOf(Integer.MAX_VALUE); + private static final BigInteger LONG_MIN = BigInteger.valueOf(Long.MIN_VALUE); + private static final BigInteger LONG_MAX = BigInteger.valueOf(Long.MAX_VALUE); + + @Override + public boolean matches(byte[] payload, int offset, int length) { + // Top-level element must be an OBJECT (low nibble == 12) to produce a row. This never collides with text + // JSON, whose first byte is '{' (0x7B) or '[' (0x5B) -- both low nibble 0x0B (ARRAY), never 0x0C. + return length >= 1 && (payload[offset] & 0x0F) == TYPE_OBJECT; + } + + @Override + public Map<String, Object> parse(byte[] payload, int offset, int length) { + int limit = offset + length; + Cursor cursor = new Cursor(payload, offset, limit); + Object value = readElement(cursor, limit); Review Comment: Confirmed and fixed in 4b7401f — good catch, this was a silent-corruption bug. Reproduced it exactly as you described: `0x0C` (OBJECT, size descriptor 0 → payload size 0) followed by real data parsed to `{}` and dropped the trailing bytes. Because `matches()` claims any payload whose first byte carries the OBJECT nibble, AUTO made it reachable for arbitrary input, so a corrupt stream message became an empty ingested row rather than a rejected one. `parse()` now enforces SQLite's exact-fill rule and reports how much was consumed: ```java if (cursor._pos != limit) { throw new IllegalArgumentException( "Top-level SQLite JSONB element consumed " + (cursor._pos - offset) + " of " + length + " bytes"); } ``` Regression tests at both levels: - `testSqliteTopLevelElementMustExactlyFillThePayload` — the bare-`0x0C`-then-data case, a well-formed `{"a":1}` with one stray trailing byte, and that a lone `0x0C` (legitimately empty object, consumes the whole payload) still parses to `{}`. - `testSqliteTrailingBytesAreRejectedRatherThanIngestedAsAPartialRow` — end-to-end through `JSONMessageDecoder`, via both AUTO and a pinned `SQLITE_JSONB`, asserting the message fails instead of yielding a partial row. ########## pinot-plugins/pinot-input-format/pinot-json/src/main/java/org/apache/pinot/plugin/inputformat/json/JSONMessageDecoder.java: ########## @@ -20,34 +20,50 @@ import java.util.Map; import java.util.Set; +import org.apache.pinot.plugin.inputformat.json.format.JsonPayloadFormat; +import org.apache.pinot.plugin.inputformat.json.format.JsonPayloadParser; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.data.readers.RecordExtractor; import org.apache.pinot.spi.plugin.PluginManager; import org.apache.pinot.spi.stream.StreamMessageDecoder; -import org.apache.pinot.spi.utils.JsonUtils; /** * An implementation of StreamMessageDecoder to read JSON records from a stream. + * + * <p>Set the {@value #JSON_FORMAT_CONFIG_KEY} decoder property to pin the payload encoding to one of + * {@code TEXT}, {@code POSTGRES_JSONB}, {@code SQLITE_JSONB}, {@code SMILE} or {@code CBOR}. + * + * <p>When unset (equivalently {@code AUTO}) the encoding is detected per message from its leading magic / + * version bytes, falling back to text JSON. Detection is allocation-free and cannot mis-route a well-formed + * text JSON document: a top-level {@code {} / {@code [} (optionally after whitespace) collides with none of the + * binary signatures. Pin {@code TEXT} to skip detection entirely. See {@link JsonPayloadFormat}. Review Comment: Fixed in 4b7401f. You're right that `{@code {}` terminates the inline tag at the first `}`. Switched to HTML entities so it renders and passes doclint: ``` * text JSON document: a top-level <code>{</code> or <code>[</code> (optionally after whitespace) collides ``` Verified with `mvn javadoc:javadoc` on the module — no errors. ########## pinot-plugins/pinot-input-format/pinot-json/src/main/java/org/apache/pinot/plugin/inputformat/json/format/JsonPayloadFormat.java: ########## @@ -0,0 +1,82 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.plugin.inputformat.json.format; + +import java.util.Locale; +import javax.annotation.Nullable; + + +/// The payload encodings the JSON stream decoder understands. Selected through the `jsonFormat` decoder +/// property; [#AUTO] (the default) sniffs each message's leading bytes. +public enum JsonPayloadFormat { + /// Detect the encoding per message from its leading magic / version bytes, falling back to text JSON. + AUTO(new AutoDetectPayloadParser()), + /// UTF-8 text JSON (the historical default). + TEXT(new TextJsonPayloadParser()), + /// PostgreSQL `jsonb` binary wire format (version byte + text JSON). + POSTGRES_JSONB(new PostgresJsonbPayloadParser()), + /// SQLite 3.45+ JSONB binary format. + SQLITE_JSONB(new SqliteJsonbPayloadParser()), + /// Jackson Smile binary JSON. + SMILE(new SmileJsonPayloadParser()), + /// CBOR (RFC 8949). + CBOR(new CborJsonPayloadParser()); + + // AUTO detection precedence: formats with strong, unambiguous magic bytes first; text (the catch-all, + // detected only by a leading '{' / '[') last so it never shadows a binary format. POSTGRES_JSONB precedes + // TEXT because its body *is* text JSON behind a version byte. + private static final JsonPayloadFormat[] DETECTION_ORDER = + {SMILE, CBOR, POSTGRES_JSONB, SQLITE_JSONB, TEXT}; + + private final JsonPayloadParser _parser; + + JsonPayloadFormat(JsonPayloadParser parser) { + _parser = parser; + } + + /// The parser for this format. For [#AUTO] this is a parser that resolves the concrete format per message. + public JsonPayloadParser getParser() { + return _parser; + } + + /// Resolves a configured format name (case-insensitive), returning [#AUTO] when unset. + public static JsonPayloadFormat fromConfig(@Nullable String value) { + if (value == null || value.trim().isEmpty()) { + return AUTO; + } + try { + return valueOf(value.trim().toUpperCase(Locale.ROOT)); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException( + "Unsupported jsonFormat '" + value + "'. Supported values: AUTO, TEXT, POSTGRES_JSONB, SQLITE_JSONB, " + + "SMILE, CBOR"); + } Review Comment: Fixed in 4b7401f — now `throw new IllegalArgumentException("Unsupported jsonFormat ...", e)`, preserving the `valueOf` failure as the cause. ########## pinot-plugins/pinot-input-format/pinot-json/src/test/java/org/apache/pinot/plugin/inputformat/json/JSONMessageDecoderBinaryTest.java: ########## @@ -0,0 +1,172 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.plugin.inputformat.json; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.dataformat.cbor.CBORFactory; +import com.fasterxml.jackson.dataformat.cbor.CBORGenerator; +import com.fasterxml.jackson.dataformat.smile.SmileFactory; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; + + +/// End-to-end coverage of {@link JSONMessageDecoder} decoding each configured / auto-detected payload format +/// through the shared {@link JSONRecordExtractor} into a {@link GenericRow}. +public class JSONMessageDecoderBinaryTest { + + private static final Set<String> RICH_FIELDS = Set.of("name", "count", "ratio"); + private static final Set<String> SINGLE_FIELD = Set.of("a"); + + private static byte[] smile(Map<String, Object> value) + throws Exception { + return new ObjectMapper(new SmileFactory()).writeValueAsBytes(value); + } + + private static byte[] cbor(Map<String, Object> value) + throws Exception { + CBORFactory factory = new CBORFactory(); + factory.enable(CBORGenerator.Feature.WRITE_TYPE_HEADER); + return new ObjectMapper(factory).writeValueAsBytes(value); + } + + private static byte[] bytes(int... values) { + byte[] result = new byte[values.length]; + for (int i = 0; i < values.length; i++) { + result[i] = (byte) values[i]; + } + return result; + } + + // Hand-built {"a": 1} fixtures matching the parser unit tests. + // SQLite: object(size 4){ text("a"), int("1") }. + private static final byte[] SQLITE_A1 = bytes(0x4C, 0x17, 0x61, 0x13, 0x31); + // PostgreSQL jsonb_send framing: version byte 1 followed by the text JSON body. + private static final byte[] POSTGRES_A1 = + bytes(0x01, '{', '"', 'a', '"', ':', '1', '}'); + + private GenericRow decode(Map<String, String> props, Set<String> fields, byte[] payload) + throws Exception { + JSONMessageDecoder decoder = new JSONMessageDecoder(); + decoder.init(props, fields, "topic"); + return decoder.decode(payload, new GenericRow()); + } + + private void assertRich(GenericRow row) { + assertEquals(row.getValue("name"), "pinot"); + assertEquals(row.getValue("count"), 7); + assertEquals(row.getValue("ratio"), 2.5); + } + + private Map<String, Object> richDoc() { + return Map.of("name", "pinot", "count", 7, "ratio", 2.5); + } + + @Test + public void testDefaultIsText() + throws Exception { + assertRich(decode(Map.of(), RICH_FIELDS, "{\"name\":\"pinot\",\"count\":7,\"ratio\":2.5}".getBytes())); + } Review Comment: Fixed in 4b7401f. Renamed to `testUnsetFormatAutoDetectsText`, with a doc comment making the intent explicit: the default is AUTO (per-message detection), and text JSON must still decode through it. ########## pinot-plugins/pinot-input-format/pinot-json/src/test/java/org/apache/pinot/plugin/inputformat/json/format/JsonPayloadFormatTest.java: ########## @@ -0,0 +1,422 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.plugin.inputformat.json.format; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.dataformat.cbor.CBORFactory; +import com.fasterxml.jackson.dataformat.cbor.CBORGenerator; +import com.fasterxml.jackson.dataformat.smile.SmileFactory; +import java.math.BigInteger; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +public class JsonPayloadFormatTest { + + private static byte[] bytes(int... values) { + byte[] result = new byte[values.length]; + for (int i = 0; i < values.length; i++) { + result[i] = (byte) values[i]; + } + return result; + } + + /// PostgreSQL `jsonb_send` framing: version byte 1 followed by the text JSON body. + private static byte[] postgres(String json) { + byte[] body = json.getBytes(StandardCharsets.UTF_8); + byte[] result = new byte[body.length + 1]; + result[0] = 1; + System.arraycopy(body, 0, result, 1, body.length); + return result; + } + + private static byte[] smile(Object value) + throws Exception { + return new ObjectMapper(new SmileFactory()).writeValueAsBytes(value); + } + + private static byte[] cborSelfDescribed(Object value) + throws Exception { + CBORFactory factory = new CBORFactory(); + factory.enable(CBORGenerator.Feature.WRITE_TYPE_HEADER); + return new ObjectMapper(factory).writeValueAsBytes(value); + } + + private static Map<String, Object> parse(JsonPayloadParser parser, byte[] payload) + throws Exception { + return parser.parse(payload, 0, payload.length); + } + + // ------------------------------------------------------------------------------------------------------ + // Config parsing + // ------------------------------------------------------------------------------------------------------ + + @Test + public void testFromConfig() { + assertEquals(JsonPayloadFormat.fromConfig(null), JsonPayloadFormat.AUTO); + assertEquals(JsonPayloadFormat.fromConfig(""), JsonPayloadFormat.AUTO); + assertEquals(JsonPayloadFormat.fromConfig(" "), JsonPayloadFormat.AUTO); + assertEquals(JsonPayloadFormat.fromConfig("text"), JsonPayloadFormat.TEXT); + assertEquals(JsonPayloadFormat.fromConfig(" Smile "), JsonPayloadFormat.SMILE); + assertEquals(JsonPayloadFormat.fromConfig("POSTGRES_JSONB"), JsonPayloadFormat.POSTGRES_JSONB); + // AUTO resolves the concrete format per message, so every format has a non-null parser. + assertSame(JsonPayloadFormat.AUTO.getParser().getClass(), AutoDetectPayloadParser.class); + assertThrows(IllegalArgumentException.class, () -> JsonPayloadFormat.fromConfig("bson")); + } + + // ------------------------------------------------------------------------------------------------------ + // Text + // ------------------------------------------------------------------------------------------------------ + + @Test + public void testText() + throws Exception { + JsonPayloadParser parser = JsonPayloadFormat.TEXT.getParser(); + Map<String, Object> map = parse(parser, " {\"a\": 1, \"b\": \"x\"}".getBytes()); + assertEquals(map.get("a"), 1); + assertEquals(map.get("b"), "x"); + assertTrue(parser.matches("{\"a\":1}".getBytes(), 0, 7)); + assertTrue(parser.matches(" [1,2]".getBytes(), 0, 7)); Review Comment: Fixed in 4b7401f. Added a `utf8(String)` helper and routed every fixture through it; `postgres(...)` now builds on it too. Did the same in `JSONMessageDecoderBinaryTest` via a `TEXT_DOC` constant using `StandardCharsets.UTF_8`. (Left the pre-existing `line.getBytes()` in `JSONMessageDecoderTest` alone — untouched by this PR.) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
