Eliminate lambda indirection overhead in DynamoDB Enhanced Client

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/mapper/ResolvedImmutableAttribute.java

@@ -27,33 +27,40 @@
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 
 @SdkInternalApi
-public final class ResolvedImmutableAttribute<T, B> {
+public final class ResolvedImmutableAttribute<T, B, R> {
     private final String attributeName;
-    private final Function<T, AttributeValue> getAttributeMethod;
     private final BiConsumer<B, AttributeValue> updateBuilderMethod;
     private final StaticTableMetadata tableMetadata;
     private final AttributeConverter attributeConverter;
+    private final Function<T, R> getter;
+    private final AttributeType<R> attributeType;
+
 
     private ResolvedImmutableAttribute(String attributeName,
-                                       Function<T, AttributeValue> getAttributeMethod,
                                        BiConsumer<B, AttributeValue> updateBuilderMethod,
                                        StaticTableMetadata tableMetadata,
-                                       AttributeConverter attributeConverter) {
+                                       AttributeConverter attributeConverter,
+                                       Function<T, R> getter,
+                                       AttributeType<R> attributeType) {
         this.attributeName = attributeName;
-        this.getAttributeMethod = getAttributeMethod;
         this.updateBuilderMethod =  updateBuilderMethod;
         this.tableMetadata = tableMetadata;
         this.attributeConverter = attributeConverter;
+        this.getter = getter;
+        this.attributeType = attributeType;
+    }
+
+    AttributeValue getAttributeValue(T item) {
+        R value = getter.apply(item);
+        return value == null ? nullAttributeValue()
+                             : attributeType.objectToAttributeValue(value);
     }
 
-    public static <T, B, R> ResolvedImmutableAttribute<T, B> create(ImmutableAttribute<T, B, R> immutableAttribute,
+    public static <T, B, R> ResolvedImmutableAttribute<T, B, R> create(ImmutableAttribute<T, B, R> immutableAttribute,
                                                                     AttributeConverter<R> attributeConverter) {
 
         AttributeType<R> attributeType = StaticAttributeType.create(attributeConverter);
-        Function<T, AttributeValue> getAttributeValueWithTransform = item -> {
-            R value = immutableAttribute.getter().apply(item);
-            return value == null ? nullAttributeValue() : attributeType.objectToAttributeValue(value);
-        };
+        Function<T, R> getter = immutableAttribute.getter();
 
         // When setting a value on the java object, do not explicitly set nulls as this can cause an NPE to be thrown
         // if the target attribute type is a primitive.
@@ -78,35 +85,35 @@
             tag.modifyMetadata(immutableAttribute.name(), attributeType.attributeValueType()).accept(tableMetadataBuilder);
         });
         return new ResolvedImmutableAttribute<>(immutableAttribute.name(),
-                                                getAttributeValueWithTransform,
                                                 updateBuilderWithTransform,
                                                 tableMetadataBuilder.build(),
-                                                attributeConverter);
+                                                attributeConverter,
+                                                getter,
+                                                attributeType);
     }
 
-    public <T1, B1> ResolvedImmutableAttribute<T1, B1> transform(
+    public <T1, B1> ResolvedImmutableAttribute<T1, B1, R> transform(
         Function<T1, T> transformItem,
         Function<B1, B> transformBuilder) {
 
         return new ResolvedImmutableAttribute<>(
             attributeName,
+            (item, value) -> updateBuilderMethod.accept(transformBuilder.apply(item), value),
+            tableMetadata,
+            attributeConverter,
             item -> {
                 T otherItem = transformItem.apply(item);
-
-                // If the containing object is null don't attempt to read attributes from it
-                return otherItem == null ?
-                    nullAttributeValue() : getAttributeMethod.apply(otherItem);
+                return otherItem == null ? null : getter.apply(otherItem);
             },
-            (item, value) -> updateBuilderMethod.accept(transformBuilder.apply(item), value),
-            tableMetadata, attributeConverter);
+            attributeType);
     }
 
     public String attributeName() {
         return attributeName;
     }
 
     public Function<T, AttributeValue> attributeGetterMethod() {
-        return getAttributeMethod;
+        return this::getAttributeValue;
     }
 
     public BiConsumer<B, AttributeValue> updateItemMethod() {

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/mapper/ImmutableAttribute.java

@@ -173,7 +173,7 @@ public Builder<T, B, R> toBuilder() {
     }
 
 
-    ResolvedImmutableAttribute<T, B> resolve(AttributeConverterProvider attributeConverterProvider) {
+    ResolvedImmutableAttribute<T, B, R> resolve(AttributeConverterProvider attributeConverterProvider) {
         return ResolvedImmutableAttribute.create(this,
                                                  converterFrom(attributeConverterProvider));
     }

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/mapper/StaticImmutableTableSchema.java

@@ -78,10 +78,10 @@
 @SdkPublicApi
 @ThreadSafe
 public final class StaticImmutableTableSchema<T, B> implements TableSchema<T> {
-    private final List<ResolvedImmutableAttribute<T, B>> attributeMappers;
+    private final List<ResolvedImmutableAttribute<T, B, ?>> attributeMappers;
     private final Supplier<B> newBuilderSupplier;
     private final Function<B, T> buildItemFunction;
-    private final Map<String, ResolvedImmutableAttribute<T, B>> indexedMappers;
+    private final Map<String, ResolvedImmutableAttribute<T, B, ?>> indexedMappers;
     private final StaticTableMetadata tableMetadata;
     private final EnhancedType<T> itemType;
     private final AttributeConverterProvider attributeConverterProvider;
@@ -203,12 +203,12 @@ private StaticImmutableTableSchema(Builder<T, B> builder) {
                 ConverterProviderResolver.resolveProviders(builder.attributeConverterProviders);
 
         // Resolve declared attributes and find converters for them
-        Stream<ResolvedImmutableAttribute<T, B>> attributesStream = builder.attributes == null ?
+        Stream<ResolvedImmutableAttribute<T, B, ?>> attributesStream = builder.attributes == null ?
             Stream.empty() : builder.attributes.stream().map(a -> a.resolve(this.attributeConverterProvider));
 
         // Merge resolved declared attributes
-        List<ResolvedImmutableAttribute<T, B>> mutableAttributeMappers = new ArrayList<>();
-        Map<String, ResolvedImmutableAttribute<T, B>>  mutableIndexedMappers = new HashMap<>();
+        List<ResolvedImmutableAttribute<T, B, ?>> mutableAttributeMappers = new ArrayList<>();
+        Map<String, ResolvedImmutableAttribute<T, B, ?>>  mutableIndexedMappers = new HashMap<>();
         Set<String> mutableAttributeNames = new LinkedHashSet<>();
         Stream.concat(attributesStream, builder.additionalAttributes.stream()).forEach(
             resolvedAttribute -> {
@@ -300,7 +300,7 @@ public static <T, B> Builder<T, B> builder(EnhancedType<T> itemType, EnhancedTyp
     public static final class Builder<T, B> {
         private final EnhancedType<T> itemType;
         private final EnhancedType<B> builderType;
-        private final List<ResolvedImmutableAttribute<T, B>> additionalAttributes = new ArrayList<>();
+        private final List<ResolvedImmutableAttribute<T, B, ?>> additionalAttributes = new ArrayList<>();
         private final List<FlattenedMapper<T, B, ?>> flattenedObjectMappers = new ArrayList<>();
 
         private FlattenedMapperForMaps<T, B> flattenedMapMapper;
@@ -447,7 +447,7 @@ public <T1> Builder<T, B> flatten(String mapName, Function<T, Map<String, String
          * table schema.
          */
         public Builder<T, B> extend(StaticImmutableTableSchema<? super T, ? super B> superTableSchema) {
-            Stream<ResolvedImmutableAttribute<T, B>> attributeStream =
+            Stream<ResolvedImmutableAttribute<T, B, ?>> attributeStream =
                 upcastingTransformForAttributes(superTableSchema.attributeMappers);
             attributeStream.forEach(this.additionalAttributes::add);
             return this;
@@ -508,8 +508,8 @@ public StaticImmutableTableSchema<T, B> build() {
             return new StaticImmutableTableSchema<>(this);
         }
 
-        private static <T extends T1, T1, B extends B1, B1> Stream<ResolvedImmutableAttribute<T, B>>
-            upcastingTransformForAttributes(Collection<ResolvedImmutableAttribute<T1, B1>> superAttributes) {
+        private static <T extends T1, T1, B extends B1, B1> Stream<ResolvedImmutableAttribute<T, B, ?>>
+            upcastingTransformForAttributes(Collection<ResolvedImmutableAttribute<T1, B1, ?>> superAttributes) {
 
             return superAttributes.stream().map(attribute -> attribute.transform(x -> x, x -> x));
         }
@@ -538,7 +538,7 @@ public T mapToItem(Map<String, AttributeValue> attributeMap, boolean preserveEmp
             AttributeValue value = entry.getValue();
 
             if (!isNullAttributeValue(value)) {
-                ResolvedImmutableAttribute<T, B> attributeMapper = indexedMappers.get(key);
+                ResolvedImmutableAttribute<T, B, ?> attributeMapper = indexedMappers.get(key);
 
                 if (attributeMapper != null) {
                     if (builder == null) {
@@ -647,7 +647,7 @@ public Map<String, AttributeValue> itemToMap(T item, Collection<String> attribut
 
     @Override
     public AttributeValue attributeValue(T item, String key) {
-        ResolvedImmutableAttribute<T, B> attributeMapper = indexedMappers.get(key);
+        ResolvedImmutableAttribute<T, B, ?> attributeMapper = indexedMappers.get(key);
 
         if (attributeMapper == null) {
             FlattenedMapper<T, B, ?> flattenedMapper = flattenedObjectMappers.get(key);
@@ -700,7 +700,7 @@ private B constructNewBuilder() {
 
     @Override
     public AttributeConverter<T> converterForAttribute(Object key) {
-        ResolvedImmutableAttribute<T, B> resolvedImmutableAttribute = indexedMappers.get(key);
+        ResolvedImmutableAttribute<T, B, ?> resolvedImmutableAttribute = indexedMappers.get(key);
         if (resolvedImmutableAttribute != null) {
             return resolvedImmutableAttribute.attributeConverter();
         }

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/mapper/ImmutableAttributeTest.java

@@ -29,6 +29,7 @@
 import org.mockito.junit.MockitoJUnitRunner;
 import software.amazon.awssdk.enhanced.dynamodb.AttributeConverter;
 import software.amazon.awssdk.enhanced.dynamodb.AttributeConverterProvider;
+import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType;
 import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
 import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.ResolvedImmutableAttribute;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -219,7 +220,7 @@
                               .attributeConverter(attributeConverter)
                               .build();
 
-        ResolvedImmutableAttribute<SimpleItem, SimpleItem> resolvedAttribute =
+        ResolvedImmutableAttribute<SimpleItem, SimpleItem, String> resolvedAttribute =
             staticAttribute.resolve(AttributeConverterProvider.defaultProvider());
 
         Function<SimpleItem, AttributeValue> attributeValueFunction = resolvedAttribute.attributeGetterMethod();
@@ -229,4 +230,43 @@
 
         assertThat(resultAttributeValue.s()).isEqualTo("test-string-custom");
     }
+
+    @Test
+    public void getAttributeValue_getterThrowsRuntimeException_propagatesException() {
+        ImmutableAttribute<SimpleItem, SimpleItem, String> staticAttribute =
+            ImmutableAttribute.builder(SimpleItem.class, SimpleItem.class, String.class)
+                              .name("test-attribute")
+                              .getter(item -> { throw new RuntimeException("getter failed"); })
+                              .setter(SimpleItem::setAString)
+                              .build();
+
+        ResolvedImmutableAttribute<SimpleItem, SimpleItem, String> resolvedAttribute =
+            staticAttribute.resolve(AttributeConverterProvider.defaultProvider());
+
+        Function<SimpleItem, AttributeValue> attributeValueFunction = resolvedAttribute.attributeGetterMethod();
+        SimpleItem item = new SimpleItem("test");
+
+        assertThatThrownBy(() -> attributeValueFunction.apply(item))
+            .isInstanceOf(RuntimeException.class)
+            .hasMessage("getter failed");
+    }
+
+    @Test
+    public void getAttributeValue_nullFromGetter_returnsNullAttributeValue() {
+        ImmutableAttribute<SimpleItem, SimpleItem, String> staticAttribute =
+            ImmutableAttribute.builder(SimpleItem.class, SimpleItem.class, String.class)
+                              .name("test-attribute")
+                              .getter(item -> null)
+                              .setter(SimpleItem::setAString)
+                              .build();
+
+        ResolvedImmutableAttribute<SimpleItem, SimpleItem, String> resolvedAttribute =
+            staticAttribute.resolve(AttributeConverterProvider.defaultProvider());
+
+        Function<SimpleItem, AttributeValue> attributeValueFunction = resolvedAttribute.attributeGetterMethod();
+        SimpleItem item = new SimpleItem("test");
+        AttributeValue result = attributeValueFunction.apply(item);
+
+        assertThat(result.nul()).isTrue();
+    }
 }