[GitHub] [iceberg] nastra commented on a change in pull request #2720: Use bulk decryption interface in ArrowReader

[GitBox <gi...@apache.org>]

nastra commented on a change in pull request #2720:
URL: https://github.com/apache/iceberg/pull/2720#discussion_r656801373



##########
File path: arrow/src/main/java/org/apache/iceberg/arrow/vectorized/ArrowReader.java
##########
@@ -210,11 +212,21 @@ public void close() throws IOException {
           .flatMap(Collection::stream)
           .collect(Collectors.toList());
       this.fileItr = fileTasks.iterator();
-      this.inputFiles = Collections.unmodifiableMap(fileTasks.stream()
+
+      Map<String, ByteBuffer> keyMetadata = Maps.newHashMap();
+      fileTasks.stream()
           .flatMap(fileScanTask -> Stream.concat(Stream.of(fileScanTask.file()), fileScanTask.deletes().stream()))
-          .map(file -> EncryptedFiles.encryptedInput(io.newInputFile(file.path().toString()), file.keyMetadata()))
-          .map(encryptionManager::decrypt)
-          .collect(Collectors.toMap(InputFile::location, Function.identity())));
+          .forEach(file -> keyMetadata.put(file.path().toString(), file.keyMetadata()));
+
+      Stream<EncryptedInputFile> encrypted = keyMetadata.entrySet().stream()
+          .map(entry -> EncryptedFiles.encryptedInput(io.newInputFile(entry.getKey()), entry.getValue()));
+
+      // decrypt with the batch call to avoid multiple RPCs to a key server, if possible
+      Iterable<InputFile> decryptedFiles = encryptionManager.decrypt(encrypted::iterator);
+
+      Map<String, InputFile> files = Maps.newHashMapWithExpectedSize(fileTasks.size());
+      decryptedFiles.forEach(decrypted -> files.putIfAbsent(decrypted.location(), decrypted));
+      this.inputFiles = Collections.unmodifiableMap(files);

Review comment:
       sgtm. I'll do the same change in `DataIterator` / `BaseDataReader` to align the 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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org