Skip to content

Commit 5ed3c31

Browse files
vbabaninrozza
vbabanin
and
rozza
authoredMar 4, 2025
Optimize BSON codec lookup. (#1632)
- Use `BsonTypeCodecMap` for decoding and encoding in `BsonArray` to enable faster access. - Optimize `BsonTypeClassMap` by replacing `Map` with a plain array to eliminate redundant hash computations. JAVA-5339 --------- Co-authored-by: Ross Lawley <ross.lawley@gmail.com>
1 parent b0e19bc commit 5ed3c31

11 files changed

+314
-54
lines changed
 

‎bson/src/main/org/bson/codecs/BsonArrayCodec.java

+12-7
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.bson.codecs.configuration.CodecRegistry;
2525

2626
import static org.bson.assertions.Assertions.notNull;
27+
import static org.bson.codecs.BsonValueCodecProvider.getBsonTypeClassMap;
2728
import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
2829

2930
/**
@@ -34,16 +35,16 @@
3435
public class BsonArrayCodec implements Codec<BsonArray> {
3536

3637
private static final CodecRegistry DEFAULT_REGISTRY = fromProviders(new BsonValueCodecProvider());
37-
38-
private final CodecRegistry codecRegistry;
38+
private static final BsonTypeCodecMap DEFAULT_BSON_TYPE_CODEC_MAP = new BsonTypeCodecMap(getBsonTypeClassMap(), DEFAULT_REGISTRY);
39+
private final BsonTypeCodecMap bsonTypeCodecMap;
3940

4041
/**
4142
* Creates a new instance with a default codec registry that uses the {@link BsonValueCodecProvider}.
4243
*
4344
* @since 3.4
4445
*/
4546
public BsonArrayCodec() {
46-
this(DEFAULT_REGISTRY);
47+
this(DEFAULT_BSON_TYPE_CODEC_MAP);
4748
}
4849

4950
/**
@@ -52,7 +53,11 @@ public BsonArrayCodec() {
5253
* @param codecRegistry the codec registry
5354
*/
5455
public BsonArrayCodec(final CodecRegistry codecRegistry) {
55-
this.codecRegistry = notNull("codecRegistry", codecRegistry);
56+
this(new BsonTypeCodecMap(getBsonTypeClassMap(), codecRegistry));
57+
}
58+
59+
private BsonArrayCodec(final BsonTypeCodecMap bsonTypeCodecMap) {
60+
this.bsonTypeCodecMap = notNull("bsonTypeCodecMap", bsonTypeCodecMap);
5661
}
5762

5863
@Override
@@ -72,7 +77,7 @@ public void encode(final BsonWriter writer, final BsonArray array, final Encoder
7277
writer.writeStartArray();
7378

7479
for (BsonValue value : array) {
75-
Codec codec = codecRegistry.get(value.getClass());
80+
Codec codec = bsonTypeCodecMap.get(value.getBsonType());
7681
encoderContext.encodeWithChildContext(codec, writer, value);
7782
}
7883

@@ -93,7 +98,7 @@ public Class<BsonArray> getEncoderClass() {
9398
* @return the non-null value read from the reader
9499
*/
95100
protected BsonValue readValue(final BsonReader reader, final DecoderContext decoderContext) {
96-
return codecRegistry.get(BsonValueCodecProvider.getClassForBsonType(reader.getCurrentBsonType())).decode(reader, decoderContext);
101+
BsonType currentBsonType = reader.getCurrentBsonType();
102+
return (BsonValue) bsonTypeCodecMap.get(currentBsonType).decode(reader, decoderContext);
97103
}
98-
99104
}

‎bson/src/main/org/bson/codecs/BsonDocumentCodec.java

+4-9
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
package org.bson.codecs;
1818

1919
import org.bson.BsonDocument;
20-
import org.bson.BsonElement;
2120
import org.bson.BsonObjectId;
2221
import org.bson.BsonReader;
2322
import org.bson.BsonType;
@@ -26,8 +25,6 @@
2625
import org.bson.codecs.configuration.CodecRegistry;
2726
import org.bson.types.ObjectId;
2827

29-
import java.util.ArrayList;
30-
import java.util.List;
3128
import java.util.Map;
3229

3330
import static org.bson.assertions.Assertions.notNull;
@@ -79,17 +76,15 @@ public CodecRegistry getCodecRegistry() {
7976

8077
@Override
8178
public BsonDocument decode(final BsonReader reader, final DecoderContext decoderContext) {
82-
List<BsonElement> keyValuePairs = new ArrayList<>();
83-
79+
BsonDocument bsonDocument = new BsonDocument();
8480
reader.readStartDocument();
8581
while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
8682
String fieldName = reader.readName();
87-
keyValuePairs.add(new BsonElement(fieldName, readValue(reader, decoderContext)));
83+
bsonDocument.append(fieldName, readValue(reader, decoderContext));
8884
}
8985

9086
reader.readEndDocument();
91-
92-
return new BsonDocument(keyValuePairs);
87+
return bsonDocument;
9388
}
9489

9590
/**
@@ -135,7 +130,7 @@ private boolean skipField(final EncoderContext encoderContext, final String key)
135130

136131
@SuppressWarnings({"unchecked", "rawtypes"})
137132
private void writeValue(final BsonWriter writer, final EncoderContext encoderContext, final BsonValue value) {
138-
Codec codec = codecRegistry.get(value.getClass());
133+
Codec codec = bsonTypeCodecMap.get(value.getBsonType());
139134
encoderContext.encodeWithChildContext(codec, writer, value);
140135
}
141136

‎bson/src/main/org/bson/codecs/BsonTypeClassMap.java

+26-35
Original file line numberDiff line numberDiff line change
@@ -31,12 +31,11 @@
3131
import org.bson.types.ObjectId;
3232
import org.bson.types.Symbol;
3333

34+
import java.util.Arrays;
3435
import java.util.Collections;
3536
import java.util.Date;
36-
import java.util.HashMap;
3737
import java.util.List;
3838
import java.util.Map;
39-
import java.util.Set;
4039

4140
/**
4241
* <p>A map from a BSON types to the Class to which it should be decoded. This class is useful if, for example,
@@ -71,7 +70,7 @@
7170
*/
7271
public class BsonTypeClassMap {
7372
static final BsonTypeClassMap DEFAULT_BSON_TYPE_CLASS_MAP = new BsonTypeClassMap();
74-
private final Map<BsonType, Class<?>> map = new HashMap<>();
73+
private final Class<?>[] bsonTypeOrdinalToClassMap = new Class<?>[256];
7574

7675
/**
7776
* Construct an instance with the default mapping, but replacing the default mapping with any values contained in the given map.
@@ -81,7 +80,7 @@ public class BsonTypeClassMap {
8180
*/
8281
public BsonTypeClassMap(final Map<BsonType, Class<?>> replacementsForDefaults) {
8382
addDefaults();
84-
map.putAll(replacementsForDefaults);
83+
replacementsForDefaults.forEach((key, value) -> bsonTypeOrdinalToClassMap[key.getValue()] = value);
8584
}
8685

8786
/**
@@ -91,41 +90,37 @@ public BsonTypeClassMap() {
9190
this(Collections.emptyMap());
9291
}
9392

94-
Set<BsonType> keys() {
95-
return map.keySet();
96-
}
97-
9893
/**
9994
* Gets the Class that is mapped to the given BSON type.
10095
*
10196
* @param bsonType the BSON type
10297
* @return the Class that is mapped to the BSON type
10398
*/
10499
public Class<?> get(final BsonType bsonType) {
105-
return map.get(bsonType);
100+
return bsonTypeOrdinalToClassMap[bsonType.getValue()];
106101
}
107102

108103
private void addDefaults() {
109-
map.put(BsonType.ARRAY, List.class);
110-
map.put(BsonType.BINARY, Binary.class);
111-
map.put(BsonType.BOOLEAN, Boolean.class);
112-
map.put(BsonType.DATE_TIME, Date.class);
113-
map.put(BsonType.DB_POINTER, BsonDbPointer.class);
114-
map.put(BsonType.DOCUMENT, Document.class);
115-
map.put(BsonType.DOUBLE, Double.class);
116-
map.put(BsonType.INT32, Integer.class);
117-
map.put(BsonType.INT64, Long.class);
118-
map.put(BsonType.DECIMAL128, Decimal128.class);
119-
map.put(BsonType.MAX_KEY, MaxKey.class);
120-
map.put(BsonType.MIN_KEY, MinKey.class);
121-
map.put(BsonType.JAVASCRIPT, Code.class);
122-
map.put(BsonType.JAVASCRIPT_WITH_SCOPE, CodeWithScope.class);
123-
map.put(BsonType.OBJECT_ID, ObjectId.class);
124-
map.put(BsonType.REGULAR_EXPRESSION, BsonRegularExpression.class);
125-
map.put(BsonType.STRING, String.class);
126-
map.put(BsonType.SYMBOL, Symbol.class);
127-
map.put(BsonType.TIMESTAMP, BsonTimestamp.class);
128-
map.put(BsonType.UNDEFINED, BsonUndefined.class);
104+
bsonTypeOrdinalToClassMap[BsonType.ARRAY.getValue()] = List.class;
105+
bsonTypeOrdinalToClassMap[BsonType.BINARY.getValue()] = Binary.class;
106+
bsonTypeOrdinalToClassMap[BsonType.BOOLEAN.getValue()] = Boolean.class;
107+
bsonTypeOrdinalToClassMap[BsonType.DATE_TIME.getValue()] = Date.class;
108+
bsonTypeOrdinalToClassMap[BsonType.DB_POINTER.getValue()] = BsonDbPointer.class;
109+
bsonTypeOrdinalToClassMap[BsonType.DOCUMENT.getValue()] = Document.class;
110+
bsonTypeOrdinalToClassMap[BsonType.DOUBLE.getValue()] = Double.class;
111+
bsonTypeOrdinalToClassMap[BsonType.INT32.getValue()] = Integer.class;
112+
bsonTypeOrdinalToClassMap[BsonType.INT64.getValue()] = Long.class;
113+
bsonTypeOrdinalToClassMap[BsonType.DECIMAL128.getValue()] = Decimal128.class;
114+
bsonTypeOrdinalToClassMap[BsonType.MAX_KEY.getValue()] = MaxKey.class;
115+
bsonTypeOrdinalToClassMap[BsonType.MIN_KEY.getValue()] = MinKey.class;
116+
bsonTypeOrdinalToClassMap[BsonType.JAVASCRIPT.getValue()] = Code.class;
117+
bsonTypeOrdinalToClassMap[BsonType.JAVASCRIPT_WITH_SCOPE.getValue()] = CodeWithScope.class;
118+
bsonTypeOrdinalToClassMap[BsonType.OBJECT_ID.getValue()] = ObjectId.class;
119+
bsonTypeOrdinalToClassMap[BsonType.REGULAR_EXPRESSION.getValue()] = BsonRegularExpression.class;
120+
bsonTypeOrdinalToClassMap[BsonType.STRING.getValue()] = String.class;
121+
bsonTypeOrdinalToClassMap[BsonType.SYMBOL.getValue()] = Symbol.class;
122+
bsonTypeOrdinalToClassMap[BsonType.TIMESTAMP.getValue()] = BsonTimestamp.class;
123+
bsonTypeOrdinalToClassMap[BsonType.UNDEFINED.getValue()] = BsonUndefined.class;
129124
}
130125

131126
@Override
@@ -139,15 +134,11 @@ public boolean equals(final Object o) {
139134

140135
BsonTypeClassMap that = (BsonTypeClassMap) o;
141136

142-
if (!map.equals(that.map)) {
143-
return false;
144-
}
145-
146-
return true;
137+
return Arrays.equals(bsonTypeOrdinalToClassMap, that.bsonTypeOrdinalToClassMap);
147138
}
148139

149140
@Override
150141
public int hashCode() {
151-
return map.hashCode();
142+
return Arrays.hashCode(bsonTypeOrdinalToClassMap);
152143
}
153144
}

‎bson/src/main/org/bson/codecs/BsonTypeCodecMap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ public class BsonTypeCodecMap {
4040
public BsonTypeCodecMap(final BsonTypeClassMap bsonTypeClassMap, final CodecRegistry codecRegistry) {
4141
this.bsonTypeClassMap = notNull("bsonTypeClassMap", bsonTypeClassMap);
4242
notNull("codecRegistry", codecRegistry);
43-
for (BsonType cur : bsonTypeClassMap.keys()) {
43+
for (BsonType cur : BsonType.values()) {
4444
Class<?> clazz = bsonTypeClassMap.get(cur);
4545
if (clazz != null) {
4646
try {

‎bson/src/main/org/bson/codecs/DocumentCodec.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -156,7 +156,7 @@ public void encode(final BsonWriter writer, final Document document, final Encod
156156

157157
beforeFields(writer, encoderContext, document);
158158

159-
for (final Map.Entry<String, Object> entry : ((Map<String, Object>) document).entrySet()) {
159+
for (final Map.Entry<String, Object> entry : document.entrySet()) {
160160
if (skipField(encoderContext, entry.getKey())) {
161161
continue;
162162
}

‎bson/src/test/unit/org/bson/codecs/IterableCodecProviderSpecification.groovy

+2-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package org.bson.codecs
1818

19+
import org.bson.BsonType
1920
import spock.lang.Specification
2021

2122
import static org.bson.codecs.configuration.CodecRegistries.fromProviders
@@ -57,7 +58,7 @@ class IterableCodecProviderSpecification extends Specification {
5758
def 'unidentical instances should not be equal'() {
5859
given:
5960
def first = new IterableCodecProvider()
60-
def second = new IterableCodecProvider(new BsonTypeClassMap([BOOLEAN: String]))
61+
def second = new IterableCodecProvider(new BsonTypeClassMap([(BsonType.BOOLEAN): String]))
6162
def third = new IterableCodecProvider(new BsonTypeClassMap(), { Object from ->
6263
from
6364
})

‎driver-benchmarks/build.gradle

+10
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,16 @@ dependencies {
3333
api project(':driver-sync')
3434
api project(':mongodb-crypt')
3535
implementation "ch.qos.logback:logback-classic:$logbackVersion"
36+
37+
implementation 'org.openjdk.jmh:jmh-core:1.37'
38+
annotationProcessor 'org.openjdk.jmh:jmh-generator-annprocess:1.37'
39+
}
40+
41+
tasks.register("jmh", JavaExec) {
42+
group = 'benchmark'
43+
description = 'Run JMH benchmarks.'
44+
mainClass = 'org.openjdk.jmh.Main'
45+
classpath = sourceSets.main.runtimeClasspath
3646
}
3747

3848
javadoc {

‎driver-benchmarks/src/main/com/mongodb/benchmark/jmh/codec/BsonArrayCodecBenchmark.java

+99
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
/*
2+
* Copyright 2016-present MongoDB, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*
16+
*/
17+
18+
package com.mongodb.benchmark.jmh.codec;
19+
20+
import com.mongodb.internal.connection.ByteBufferBsonOutput;
21+
import com.mongodb.internal.connection.PowerOfTwoBufferPool;
22+
import org.bson.BsonArray;
23+
import org.bson.BsonBinaryReader;
24+
import org.bson.BsonBinaryWriter;
25+
import org.bson.BsonDocument;
26+
import org.bson.BsonDouble;
27+
import org.bson.codecs.BsonArrayCodec;
28+
import org.bson.codecs.DecoderContext;
29+
import org.bson.codecs.EncoderContext;
30+
import org.jetbrains.annotations.NotNull;
31+
import org.openjdk.jmh.annotations.Benchmark;
32+
import org.openjdk.jmh.annotations.BenchmarkMode;
33+
import org.openjdk.jmh.annotations.Fork;
34+
import org.openjdk.jmh.annotations.Level;
35+
import org.openjdk.jmh.annotations.Measurement;
36+
import org.openjdk.jmh.annotations.Mode;
37+
import org.openjdk.jmh.annotations.OutputTimeUnit;
38+
import org.openjdk.jmh.annotations.Scope;
39+
import org.openjdk.jmh.annotations.Setup;
40+
import org.openjdk.jmh.annotations.State;
41+
import org.openjdk.jmh.annotations.Warmup;
42+
import org.openjdk.jmh.infra.Blackhole;
43+
44+
import java.io.IOException;
45+
import java.nio.ByteBuffer;
46+
import java.util.concurrent.TimeUnit;
47+
import static com.mongodb.benchmark.jmh.codec.BsonUtils.getDocumentAsBuffer;
48+
49+
@BenchmarkMode(Mode.Throughput)
50+
@Warmup(iterations = 20, time = 2, timeUnit = TimeUnit.SECONDS)
51+
@Measurement(iterations = 20, time = 2, timeUnit = TimeUnit.SECONDS)
52+
@OutputTimeUnit(TimeUnit.SECONDS)
53+
@Fork(3)
54+
public class BsonArrayCodecBenchmark {
55+
56+
@State(Scope.Benchmark)
57+
public static class Input {
58+
protected final PowerOfTwoBufferPool bufferPool = PowerOfTwoBufferPool.DEFAULT;
59+
protected final BsonArrayCodec bsonArrayCodec = new BsonArrayCodec();
60+
protected BsonDocument document;
61+
protected byte[] documentBytes;
62+
private BsonBinaryReader reader;
63+
private BsonBinaryWriter writer;
64+
private BsonArray bsonValues;
65+
66+
@Setup
67+
public void setup() throws IOException {
68+
bsonValues = new BsonArray();
69+
document = new BsonDocument("array", bsonValues);
70+
71+
for (int i = 0; i < 1000; i++) {
72+
bsonValues.add(new BsonDouble(i));
73+
}
74+
75+
documentBytes = getDocumentAsBuffer(document);
76+
}
77+
78+
@Setup(Level.Invocation)
79+
public void beforeIteration() {
80+
reader = new BsonBinaryReader(ByteBuffer.wrap(documentBytes));
81+
writer = new BsonBinaryWriter(new ByteBufferBsonOutput(bufferPool));
82+
83+
reader.readStartDocument();
84+
writer.writeStartDocument();
85+
writer.writeName("array");
86+
}
87+
}
88+
89+
@Benchmark
90+
public void decode(@NotNull Input input, @NotNull Blackhole blackhole) {
91+
blackhole.consume(input.bsonArrayCodec.decode(input.reader, DecoderContext.builder().build()));
92+
}
93+
94+
@Benchmark
95+
public void encode(@NotNull Input input, @NotNull Blackhole blackhole) {
96+
input.bsonArrayCodec.encode(input.writer, input.bsonValues, EncoderContext.builder().build());
97+
blackhole.consume(input);
98+
}
99+
}

0 commit comments

Comments
 (0)
Failed to load comments.