xiangfu0 commented on code in PR #18950: URL: https://github.com/apache/pinot/pull/18950#discussion_r3555518075
########## pinot-plugins/pinot-input-format/pinot-bson/src/main/java/org/apache/pinot/plugin/inputformat/bson/BSONRecordReader.java: ########## @@ -0,0 +1,157 @@ +/** + * 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; + // A read error hit while fetching the next document. Deferred so the current record is still emitted; it + // surfaces on the following next() call rather than discarding an already-read valid record. + private IOException _fetchError; + + 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); + _fetchError = null; + _nextDocument = readNextDocument(); + } + + @Override + public boolean hasNext() { + return _nextDocument != null || _fetchError != null; + } + + @Override + public GenericRow next(GenericRow reuse) + throws IOException { + if (_fetchError != null) { + IOException error = _fetchError; + _fetchError = null; + throw new RecordFetchException("Failed to read next BSON record", error); + } + byte[] documentBytes = _nextDocument; + Document document; + try { + document = BSONUtils.decodeDocument(documentBytes); + } catch (RuntimeException e) { + // Corrupt frame: advance past it (bytes already consumed) so we don't retry, then report a parse error. + advance(); + throw new IOException("Failed to decode BSON record", e); + } + _recordExtractor.extract(document, reuse); + advance(); + return reuse; + } + + /// Advances the look-ahead to the next framed document. A read error is stashed rather than thrown so the + /// record just returned by next() is still emitted; the error surfaces on the following next() call. + private void advance() { + try { + _nextDocument = readNextDocument(); + } catch (IOException e) { + _nextDocument = null; + _fetchError = e; + } + } + + /// Reads the next framed BSON document in full, or returns `null` at a clean end-of-stream. Throws when the + /// stream ends partway through a document, or the length prefix is invalid. + @Nullable + private byte[] readNextDocument() + throws IOException { + int b0 = _inputStream.read(); + if (b0 == -1) { + return null; + } + int b1 = _inputStream.read(); + int b2 = _inputStream.read(); + int b3 = _inputStream.read(); + if ((b1 | b2 | b3) < 0) { + throw new IOException("Truncated BSON document: incomplete length prefix"); + } + // BSON length prefix is a little-endian int32 inclusive of these 4 bytes. + int length = (b0 & 0xFF) | ((b1 & 0xFF) << 8) | ((b2 & 0xFF) << 16) | ((b3 & 0xFF) << 24); + if (length < MIN_DOCUMENT_LENGTH) { + throw new IOException("Invalid BSON document length: " + length); + } + byte[] document = new byte[length]; Review Comment: Fixed in dfe5b7fe4f — added a `MAX_DOCUMENT_LENGTH` (16MB, the BSON max) upper bound so an oversized/corrupt length prefix is rejected with an `IOException` before `new byte[length]` runs; you're right that the `OutOfMemoryError` was an `Error` and escaped both recovery paths. Regression test `testLengthAboveMaximumThrowsInsteadOfAllocating` uses your `FF FF FF 7F` example. Also closed the `_inputStream` leak in `open()` per your nit. _🤖 Addressed by [Claude Code](https://claude.com/claude-code)_ ########## pinot-plugins/pinot-input-format/pinot-bson/src/main/java/org/apache/pinot/plugin/inputformat/bson/BSONRecordExtractor.java: ########## @@ -0,0 +1,116 @@ +/** + * 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 com.google.common.collect.Maps; +import java.sql.Timestamp; +import java.util.Date; +import java.util.List; +import java.util.Map; +import org.apache.pinot.spi.data.readers.BaseRecordExtractor; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.bson.types.Binary; +import org.bson.types.Decimal128; +import org.bson.types.ObjectId; + + +/// Extracts a Pinot [GenericRow] from a decoded BSON document (`org.bson.Document`, which is a +/// `Map<String, Object>`). Values are the Java objects produced by the standard MongoDB +/// [org.bson.codecs.DocumentCodec]. +/// +/// **BSON type → Java output type:** +/// - `Double` / `Int32` / `Int64` / `Boolean` / `String` → same boxed type (pass-through) +/// - `Document` (embedded) → `Map<String, Object>` (values recursively converted) +/// - `Array` → `Object[]` (elements recursively converted) +/// - `ObjectId` → `String` (24-char hex) +/// - `DateTime` → `java.sql.Timestamp` +/// - `Decimal128` → `BigDecimal` (`NaN` / `Infinity` → `null`, as `BigDecimal` cannot represent them) +/// - `Binary` → `byte[]` +/// - `null` → `null` +/// +/// Any other (rare, deprecated, or internal) BSON type falls back to `value.toString()`. The converted values +/// follow the shared `RecordExtractor` contract, so the downstream data-type transformer coerces them to the +/// declared column type. +public class BSONRecordExtractor extends BaseRecordExtractor<Map<String, Object>> { + + @Override + public GenericRow extract(Map<String, Object> from, GenericRow to) { + if (_extractAll) { + for (Map.Entry<String, Object> entry : from.entrySet()) { + Object value = entry.getValue(); + to.putValue(entry.getKey(), value != null ? convert(value) : null); + } + } else { + for (String fieldName : _fields) { + Object value = from.get(fieldName); + to.putValue(fieldName, value != null ? convert(value) : null); + } + } + return to; + } + + @SuppressWarnings("unchecked") + private static Object convert(Object value) { + if (value instanceof Map) { + return convertMap((Map<String, Object>) value); + } + if (value instanceof List) { + return convertList((List<Object>) value); + } + if (value instanceof ObjectId) { + return ((ObjectId) value).toHexString(); + } + if (value instanceof Date) { + return new Timestamp(((Date) value).getTime()); + } + if (value instanceof Decimal128) { + Decimal128 decimal128 = (Decimal128) value; + // NaN / Infinity are legal Decimal128 values with no BigDecimal representation; surface them as null + // instead of letting bigDecimalValue() throw. + return decimal128.isNaN() || decimal128.isInfinite() ? null : decimal128.bigDecimalValue(); Review Comment: Fixed in dfe5b7fe4f — confirmed negative zero has `isNaN()==false`/`isInfinite()==false` and throws, so it now converts to `BigDecimal.ZERO`. Note there's no public negative-zero predicate and `Decimal128.parse("-0.00").equals(NEGATIVE_ZERO)` is `false` (it's negative-zero at any exponent), so I catch the `ArithmeticException` rather than pre-checking; NaN/Infinity are still short-circuited to `null` before that. Regression test covers both `-0` and `-0.00`. _🤖 Addressed by [Claude Code](https://claude.com/claude-code)_ -- 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]
