Integrate analytics accelerator into iceberg-aws

Author: SanjayMarreddi

aws/src/integration/java/org/apache/iceberg/aws/s3/TestS3FileIOIntegration.java

@@ -68,6 +68,7 @@
 import software.amazon.awssdk.services.s3control.S3ControlClient;
 import software.amazon.awssdk.utils.ImmutableMap;
 import software.amazon.awssdk.utils.IoUtils;
+import software.amazon.s3.analyticsaccelerator.util.PrefetchMode;
 
 public class TestS3FileIOIntegration {
 
@@ -255,6 +256,33 @@ public void testNewInputStreamWithMultiRegionAccessPoint() throws Exception {
     validateRead(s3FileIO);
   }
 
+  @Test
+  public void testNewInputStreamWithAnalyticsAccelerator() throws Exception {
+    s3.putObject(
+        PutObjectRequest.builder().bucket(bucketName).key(objectKey).build(),
+        RequestBody.fromBytes(contentBytes));
+    S3FileIO s3FileIO = new S3FileIO(clientFactory::s3);
+    s3FileIO.initialize(
+        ImmutableMap.of(S3FileIOProperties.S3_ANALYTICS_ACCELERATOR_ENABLED, String.valueOf(true)));
+    validateRead(s3FileIO);
+  }
+
+  @Test
+  public void testNewInputStreamWithAnalyticsAcceleratorCustomConfigured() throws Exception {
+    final String prefetchingMode = "logicalio.prefetching.mode";
+    final String s3Uri = String.format("s3://%s/%s/%s", bucketName, prefix, "someFile.parquet");
+    S3FileIO s3FileIO = new S3FileIO(clientFactory::s3);
+    s3FileIO.initialize(
+        ImmutableMap.of(
+            S3FileIOProperties.S3_ANALYTICS_ACCELERATOR_ENABLED,
+            String.valueOf(true),
+            S3FileIOProperties.S3_ANALYTICS_ACCELERATOR_PREFIX + prefetchingMode,
+            PrefetchMode.ALL.name()));
+    write(s3FileIO, s3Uri);
+    validateRead(s3FileIO, s3Uri);
+    s3FileIO.deleteFile(s3Uri);
+  }
+
   @Test
   public void testNewOutputStream() throws Exception {
     S3FileIO s3FileIO = new S3FileIO(clientFactory::s3);
@@ -324,6 +352,19 @@ public void testNewOutputStreamWithMultiRegionAccessPoint() throws Exception {
     }
   }
 
+  @Test
+  public void testNewOutputStreamWithAnalyticsAccelerator() throws Exception {
+    S3FileIO s3FileIO = new S3FileIO(clientFactory::s3);
+    s3FileIO.initialize(
+        ImmutableMap.of(S3FileIOProperties.S3_ANALYTICS_ACCELERATOR_ENABLED, String.valueOf(true)));
+    write(s3FileIO);
+    try (InputStream stream =
+        s3.getObject(GetObjectRequest.builder().bucket(bucketName).key(objectKey).build())) {
+      String result = IoUtils.toUtf8String(stream);
+      assertThat(result).isEqualTo(content);
+    }
+  }
+
   @Test
   public void testServerSideS3Encryption() throws Exception {
     S3FileIOProperties properties = new S3FileIOProperties();

aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java

@@ -28,7 +28,9 @@
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.glue.GlueClient;
 import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.S3CrtAsyncClientBuilder;
 import software.amazon.awssdk.services.sts.StsClient;
 import software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
 import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
@@ -52,6 +54,14 @@ public S3Client s3() {
         .build();
   }
 
+  @Override
+  public S3AsyncClient s3Async() {
+    return S3AsyncClient.crtBuilder()
+        .applyMutation(this::applyAssumeRoleConfigurations)
+        .applyMutation(awsClientProperties::applyAsyncConfigurations)
+        .build();
+  }
+
   @Override
   public GlueClient glue() {
     return GlueClient.builder()
@@ -95,24 +105,19 @@ public void initialize(Map<String, String> properties) {
 
   protected <T extends AwsClientBuilder & AwsSyncClientBuilder> T applyAssumeRoleConfigurations(
       T clientBuilder) {
-    AssumeRoleRequest assumeRoleRequest =
-        AssumeRoleRequest.builder()
-            .roleArn(awsProperties.clientAssumeRoleArn())
-            .roleSessionName(roleSessionName)
-            .durationSeconds(awsProperties.clientAssumeRoleTimeoutSec())
-            .externalId(awsProperties.clientAssumeRoleExternalId())
-            .tags(awsProperties.stsClientAssumeRoleTags())
-            .build();
     clientBuilder
-        .credentialsProvider(
-            StsAssumeRoleCredentialsProvider.builder()
-                .stsClient(sts())
-                .refreshRequest(assumeRoleRequest)
-                .build())
+        .credentialsProvider(createCredentialsProvider())
         .region(Region.of(awsProperties.clientAssumeRoleRegion()));
     return clientBuilder;
   }
 
+  protected S3CrtAsyncClientBuilder applyAssumeRoleConfigurations(
+      S3CrtAsyncClientBuilder clientBuilder) {
+    return clientBuilder
+        .credentialsProvider(createCredentialsProvider())
+        .region(Region.of(awsProperties.clientAssumeRoleRegion()));
+  }
+
   protected String region() {
     return awsProperties.clientAssumeRoleRegion();
   }
@@ -145,4 +150,21 @@ private String genSessionName() {
     }
     return String.format("iceberg-aws-%s", UUID.randomUUID());
   }
+
+  private StsAssumeRoleCredentialsProvider createCredentialsProvider() {
+    return StsAssumeRoleCredentialsProvider.builder()
+        .stsClient(sts())
+        .refreshRequest(createAssumeRoleRequest())
+        .build();
+  }
+
+  private AssumeRoleRequest createAssumeRoleRequest() {
+    return AssumeRoleRequest.builder()
+        .roleArn(awsProperties.clientAssumeRoleArn())
+        .roleSessionName(roleSessionName)
+        .durationSeconds(awsProperties.clientAssumeRoleTimeoutSec())
+        .externalId(awsProperties.clientAssumeRoleExternalId())
+        .tags(awsProperties.stsClientAssumeRoleTags())
+        .build();
+  }
 }

aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java

@@ -39,6 +39,7 @@
 import software.amazon.awssdk.services.glue.GlueClient;
 import software.amazon.awssdk.services.glue.GlueClientBuilder;
 import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.S3ClientBuilder;
 import software.amazon.awssdk.services.s3.S3Configuration;
@@ -118,6 +119,13 @@ public S3Client s3() {
           .build();
     }
 
+    @Override
+    public S3AsyncClient s3Async() {
+      return S3AsyncClient.crtBuilder()
+          .applyMutation(awsClientProperties::applyAsyncConfigurations)
+          .build();
+    }
+
     @Override
     public GlueClient glue() {
       return GlueClient.builder()

aws/src/main/java/org/apache/iceberg/aws/AwsClientFactory.java

@@ -23,6 +23,7 @@
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.glue.GlueClient;
 import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.S3Client;
 
 /**
@@ -38,6 +39,13 @@ public interface AwsClientFactory extends Serializable {
    */
   S3Client s3();
 
+  /**
+   * create a Amazon S3 async client
+   *
+   * @return s3 async client
+   */
+  S3AsyncClient s3Async();
+
   /**
    * create a AWS Glue client
    *

aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java

@@ -40,6 +40,7 @@
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3CrtAsyncClientBuilder;
 
 public class AwsClientProperties implements Serializable {
   /**
@@ -74,6 +75,14 @@ public class AwsClientProperties implements Serializable {
    */
   public static final String CLIENT_REGION = "client.region";
 
+  /**
+   * Used to set the maxConcurrency field in the {@link S3CrtAsyncClientBuilder} when building async
+   * s3 clients.
+   */
+  public static final String CLIENT_MAX_CONCURRENCY = "client.max-concurrency";
+
+  public static final int CLIENT_MAX_CONCURRENCY_DEFAULT = 500;
+
   /**
    * When set, the {@link VendedCredentialsProvider} will be used to fetch and refresh vended
    * credentials from this endpoint.
@@ -84,6 +93,7 @@ public class AwsClientProperties implements Serializable {
   public static final String REFRESH_CREDENTIALS_ENABLED = "client.refresh-credentials-enabled";
 
   private String clientRegion;
+  private final Integer clientMaxConcurrency;
   private final String clientCredentialsProvider;
   private final Map<String, String> clientCredentialsProviderProperties;
   private final String refreshCredentialsEndpoint;
@@ -92,6 +102,7 @@ public class AwsClientProperties implements Serializable {
 
   public AwsClientProperties() {
     this.clientRegion = null;
+    this.clientMaxConcurrency = CLIENT_MAX_CONCURRENCY_DEFAULT;
     this.clientCredentialsProvider = null;
     this.clientCredentialsProviderProperties = null;
     this.refreshCredentialsEndpoint = null;
@@ -102,6 +113,9 @@ public AwsClientProperties() {
   public AwsClientProperties(Map<String, String> properties) {
     this.allProperties = SerializableMap.copyOf(properties);
     this.clientRegion = properties.get(CLIENT_REGION);
+    this.clientMaxConcurrency =
+        PropertyUtil.propertyAsInt(
+            properties, CLIENT_MAX_CONCURRENCY, CLIENT_MAX_CONCURRENCY_DEFAULT);
     this.clientCredentialsProvider = properties.get(CLIENT_CREDENTIALS_PROVIDER);
     this.clientCredentialsProviderProperties =
         PropertyUtil.propertiesWithPrefix(properties, CLIENT_CREDENTIAL_PROVIDER_PREFIX);
@@ -135,6 +149,19 @@ public <T extends AwsClientBuilder> void applyClientRegionConfiguration(T builde
     }
   }
 
+  /**
+   * Configure async client settings.
+   *
+   * <p>Sample usage:
+   *
+   * <pre>
+   *     S3AsyncClient.crtBuilder().applyMutation(awsClientProperties::applyAsyncConfigurations)
+   * </pre>
+   */
+  public <T extends S3CrtAsyncClientBuilder> void applyAsyncConfigurations(T builder) {
+    builder.maxConcurrency(clientMaxConcurrency);
+  }
+
   /**
    * Configure the credential provider for AWS clients.
    *

aws/src/main/java/org/apache/iceberg/aws/s3/BaseS3File.java

@@ -27,14 +27,20 @@
 
 abstract class BaseS3File {
   private final S3Client client;
+  private final S3InputStreamFactory inputStreamFactory;
   private final S3URI uri;
   private final S3FileIOProperties s3FileIOProperties;
   private HeadObjectResponse metadata;
   private final MetricsContext metrics;
 
   BaseS3File(
-      S3Client client, S3URI uri, S3FileIOProperties s3FileIOProperties, MetricsContext metrics) {
+      S3Client client,
+      S3InputStreamFactory inputStreamFactory,
+      S3URI uri,
+      S3FileIOProperties s3FileIOProperties,
+      MetricsContext metrics) {
     this.client = client;
+    this.inputStreamFactory = inputStreamFactory;
     this.uri = uri;
     this.s3FileIOProperties = s3FileIOProperties;
     this.metrics = metrics;
@@ -48,6 +54,10 @@ S3Client client() {
     return client;
   }
 
+  S3InputStreamFactory inputStreamFactory() {
+    return inputStreamFactory;
+  }
+
   S3URI uri() {
     return uri;
   }

aws/src/main/java/org/apache/iceberg/aws/s3/DefaultS3FileIOAwsClientFactory.java

@@ -21,14 +21,15 @@
 import java.util.Map;
 import org.apache.iceberg.aws.AwsClientProperties;
 import org.apache.iceberg.aws.HttpClientProperties;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.S3Client;
 
-class DefaultS3FileIOAwsClientFactory implements S3FileIOAwsClientFactory {
+public class DefaultS3FileIOAwsClientFactory implements S3FileIOAwsClientFactory {
   private S3FileIOProperties s3FileIOProperties;
   private HttpClientProperties httpClientProperties;
   private AwsClientProperties awsClientProperties;
 
-  DefaultS3FileIOAwsClientFactory() {
+  public DefaultS3FileIOAwsClientFactory() {
     this.s3FileIOProperties = new S3FileIOProperties();
     this.httpClientProperties = new HttpClientProperties();
     this.awsClientProperties = new AwsClientProperties();
@@ -58,4 +59,11 @@ public S3Client s3() {
         .applyMutation(s3FileIOProperties::applyRetryConfigurations)
         .build();
   }
+
+  @Override
+  public S3AsyncClient s3Async() {
+    return S3AsyncClient.crtBuilder()
+        .applyMutation(awsClientProperties::applyAsyncConfigurations)
+        .build();
+  }
 }

aws/src/main/java/org/apache/iceberg/aws/s3/DefaultS3InputStreamFactory.java

@@ -0,0 +1,36 @@
+/*
+ * 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.iceberg.aws.s3;
+
+import org.apache.iceberg.io.SeekableInputStream;
+import org.apache.iceberg.metrics.MetricsContext;
+import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.s3.analyticsaccelerator.request.ObjectMetadata;
+
+public class DefaultS3InputStreamFactory implements S3InputStreamFactory {
+  @Override
+  public SeekableInputStream createStream(
+      S3Client client,
+      S3URI uri,
+      S3FileIOProperties properties,
+      MetricsContext metrics,
+      ObjectMetadata metadata) {
+    return new S3InputStream(client, uri, properties, metrics);
+  }
+}