xiangfu0 commented on code in PR #18950: URL: https://github.com/apache/pinot/pull/18950#discussion_r3555176787
########## pinot-plugins/pinot-input-format/pinot-bson/src/main/java/org/apache/pinot/plugin/inputformat/bson/BSONRecordReader.java: ########## @@ -0,0 +1,142 @@ +/** + * 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.bson; + +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.data.readers.RecordFetchException; +import org.apache.pinot.spi.data.readers.RecordReader; +import org.apache.pinot.spi.data.readers.RecordReaderConfig; +import org.apache.pinot.spi.data.readers.RecordReaderUtils; +import org.bson.Document; + + +/** + * Record reader for a BSON file: a concatenation of framed BSON documents (the {@code mongodump} layout), each + * self-delimited by a leading little-endian int32 byte length. Documents are read sequentially, with the next + * document's bytes fetched ahead so {@link #hasNext} does not perform I/O. GZIP-compressed files are supported. + */ +public class BSONRecordReader implements RecordReader { + // Minimum size of a BSON document: 4-byte length prefix + 1-byte terminating NUL of an empty document. + private static final int MIN_DOCUMENT_LENGTH = 5; + + private File _dataFile; + private BSONRecordExtractor _recordExtractor; + private InputStream _inputStream; + // Bytes of the next framed document, or null once the stream is exhausted. + private byte[] _nextDocument; + + public BSONRecordReader() { + } + + @Override + public void init(File dataFile, @Nullable Set<String> fieldsToRead, @Nullable RecordReaderConfig recordReaderConfig) + throws IOException { + _dataFile = dataFile; + _recordExtractor = new BSONRecordExtractor(); + _recordExtractor.init(fieldsToRead, null); + open(); + } + + private void open() + throws IOException { + _inputStream = RecordReaderUtils.getBufferedInputStream(_dataFile); + _nextDocument = readNextDocument(); + } + + @Override + public boolean hasNext() { + return _nextDocument != null; + } + + @Override + public GenericRow next(GenericRow reuse) + throws IOException { + byte[] documentBytes = _nextDocument; + // Advance the look-ahead first so that, even if the current record is corrupt, the reader has already moved + // past it and the caller can skip the failure and continue. + try { + _nextDocument = readNextDocument(); + } catch (IOException e) { + _nextDocument = null; + throw new RecordFetchException("Failed to read next BSON record", e); + } Review Comment: Good catch — fixed in 1dfdac0044. The reader now emits the already-read valid record and stashes the fetch error, surfacing it on the following `next()` call (with `hasNext()` returning true while the error is pending), so a truncated/corrupt tail no longer discards the last good record. This matches `JSONRecordReader`'s emit-then-error behavior. Added a regression test (`testValidRecordEmittedBeforeTruncatedTail`). ########## pinot-plugins/pinot-input-format/pinot-bson/src/test/java/org/apache/pinot/plugin/inputformat/bson/BSONMessageDecoderTest.java: ########## @@ -0,0 +1,148 @@ +/** + * 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.bson; + +import java.math.BigDecimal; +import java.sql.Timestamp; +import java.util.Date; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.bson.BsonBinaryWriter; +import org.bson.Document; +import org.bson.codecs.DocumentCodec; +import org.bson.codecs.EncoderContext; +import org.bson.io.BasicOutputBuffer; +import org.bson.types.Decimal128; +import org.bson.types.ObjectId; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +public class BSONMessageDecoderTest { + + @Test + public void testDecodeScalarsAndLogicalTypes() + throws Exception { + ObjectId objectId = new ObjectId("5f2b1c0e1c9d440000a1b2c3"); + long epochMillis = 1_588_469_340_000L; + Document document = new Document() + .append("id", objectId) + .append("name", "widget") + .append("count", 7) + .append("views", 1_588_469_340_000L) + .append("price", 9.5d) + .append("active", true) + .append("amount", Decimal128.parse("123.456")) + .append("createdAt", new Date(epochMillis)); + + Set<String> fields = document.keySet(); + BSONMessageDecoder decoder = new BSONMessageDecoder(); + decoder.init(Map.of(), fields, "testTopic"); + + GenericRow row = new GenericRow(); + decoder.decode(encode(document), row); + + assertEquals(row.getValue("id"), objectId.toHexString()); + assertEquals(row.getValue("name"), "widget"); + assertEquals(row.getValue("count"), 7); + assertEquals(row.getValue("views"), 1_588_469_340_000L); + assertEquals(row.getValue("price"), 9.5d); + assertEquals(row.getValue("active"), true); + assertEquals(row.getValue("amount"), new BigDecimal("123.456")); + assertEquals(row.getValue("createdAt"), new Timestamp(epochMillis)); + } + + @Test + public void testDecodeNestedDocumentAndArray() + throws Exception { + Document document = new Document() + .append("tags", List.of("a", "b", "c")) + .append("meta", new Document("region", "us-east").append("shard", 3)); + + Set<String> fields = document.keySet(); + BSONMessageDecoder decoder = new BSONMessageDecoder(); + decoder.init(Map.of(), fields, "testTopic"); + + GenericRow row = new GenericRow(); + decoder.decode(encode(document), row); + + assertEquals((Object[]) row.getValue("tags"), new Object[]{"a", "b", "c"}); + Map<?, ?> meta = (Map<?, ?>) row.getValue("meta"); + assertEquals(meta.get("region"), "us-east"); + assertEquals(meta.get("shard"), 3); + } + + @Test + public void testDecodeWithFieldProjection() + throws Exception { + Document document = new Document().append("keep", 1).append("drop", 2); + + BSONMessageDecoder decoder = new BSONMessageDecoder(); + decoder.init(Map.of(), Set.of("keep"), "testTopic"); + + GenericRow row = new GenericRow(); + decoder.decode(encode(document), row); + + assertEquals(row.getValue("keep"), 1); + assertNull(row.getValue("drop")); + } + + @Test + public void testDecodeWithOffsetAndLength() + throws Exception { + Document document = new Document().append("k", "v"); + byte[] encoded = encode(document); + + // Embed the payload inside a larger buffer to exercise the offset/length decode path. + byte[] padded = new byte[encoded.length + 8]; + int offset = 5; + System.arraycopy(encoded, 0, padded, offset, encoded.length); + + BSONMessageDecoder decoder = new BSONMessageDecoder(); + decoder.init(Map.of(), Set.of("k"), "testTopic"); + + GenericRow row = new GenericRow(); + GenericRow result = decoder.decode(padded, offset, encoded.length, row); + assertTrue(result == row); Review Comment: Done in 1dfdac0044 — switched to `assertSame(result, row)`. -- 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]
