PR comments

spotify · Dec 18, 2023 · 875ae0e · 875ae0e
1 parent 1925d7e
commit 875ae0e
Show file tree

Hide file tree

Showing 9 changed files with 22 additions and 20 deletions.
diff --git a/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/AvroBucketMetadata.java b/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/AvroBucketMetadata.java
@@ -135,12 +135,12 @@ public Map<Class<?>, Coder<?>> coderOverrides() {
   }
 
   @Override
-  public int hashPrimaryKeyMetadata() {
+  int hashPrimaryKeyMetadata() {
     return Objects.hash(keyField, getKeyClass());
   }
 
   @Override
-  public int hashSecondaryKeyMetadata() {
+  int hashSecondaryKeyMetadata() {
     return Objects.hash(keyFieldSecondary, getKeyClassSecondary());
   }
 

diff --git a/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/BucketMetadata.java b/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/BucketMetadata.java
@@ -327,9 +327,9 @@ public Set<Class<? extends BucketMetadata>> compatibleMetadataTypes() {
 
   public abstract K2 extractKeySecondary(V value);
 
-  public abstract int hashPrimaryKeyMetadata();
+  abstract int hashPrimaryKeyMetadata();
 
-  public abstract int hashSecondaryKeyMetadata();
+  abstract int hashSecondaryKeyMetadata();
 
   public SortedBucketIO.ComparableKeyBytes primaryComparableKeyBytes(V value) {
     return new SortedBucketIO.ComparableKeyBytes(getKeyBytesPrimary(value), null);

diff --git a/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/JsonBucketMetadata.java b/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/JsonBucketMetadata.java
@@ -148,12 +148,12 @@ public void populateDisplayData(Builder builder) {
   }
 
   @Override
-  public int hashPrimaryKeyMetadata() {
+  int hashPrimaryKeyMetadata() {
     return Objects.hash(keyField, getKeyClass());
   }
 
   @Override
-  public int hashSecondaryKeyMetadata() {
+  int hashSecondaryKeyMetadata() {
     return Objects.hash(keyFieldSecondary, getKeyClassSecondary());
   }
 }
diff --git a/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/ParquetBucketMetadata.java b/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/ParquetBucketMetadata.java
@@ -192,12 +192,12 @@ public void populateDisplayData(DisplayData.Builder builder) {
   }
 
   @Override
-  public int hashPrimaryKeyMetadata() {
+  int hashPrimaryKeyMetadata() {
     return Objects.hash(keyField, getKeyClass());
   }
 
   @Override
-  public int hashSecondaryKeyMetadata() {
+  int hashSecondaryKeyMetadata() {
     return Objects.hash(keyFieldSecondary, getKeyClassSecondary());
   }
 

diff --git a/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/SortedBucketSource.java b/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/SortedBucketSource.java
@@ -495,7 +495,9 @@ public Predicate<V> getPredicate() {
     }
 
     public Coder<V> getCoder() {
-      return directories.entrySet().iterator().next().getValue().getValue().getCoder();
+      final KV<String, FileOperations<V>> sampledSource =
+          directories.entrySet().iterator().next().getValue();
+      return sampledSource.getValue().getCoder();
     }
 
     static CoGbkResultSchema schemaOf(List<BucketedInput<?>> sources) {
@@ -523,13 +525,13 @@ long getOrSampleByteSize() {
               entry -> {
                 // Take at most 10 buckets from the directory to sample
                 // Check for single-shard filenames template first, then multi-shard
+                final String filenameSuffix = entry.getValue().getKey();
                 List<Metadata> sampledFiles =
-                    sampleDirectory(entry.getKey(), "*-0000?-of-?????" + entry.getValue().getKey());
+                    sampleDirectory(entry.getKey(), "*-0000?-of-?????" + filenameSuffix);
                 if (sampledFiles.isEmpty()) {
                   sampledFiles =
                       sampleDirectory(
-                          entry.getKey(),
-                          "*-0000?-of-*-shard-00000-of-?????" + entry.getValue().getKey());
+                          entry.getKey(), "*-0000?-of-*-shard-00000-of-?????" + filenameSuffix);
                 }
 
                 int numBuckets = 0;

diff --git a/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/TensorFlowBucketMetadata.java b/scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/TensorFlowBucketMetadata.java
@@ -159,12 +159,12 @@ public void populateDisplayData(Builder builder) {
   }
 
   @Override
-  public int hashPrimaryKeyMetadata() {
+  int hashPrimaryKeyMetadata() {
     return Objects.hash(keyField, getKeyClass());
   }
 
   @Override
-  public int hashSecondaryKeyMetadata() {
+  int hashSecondaryKeyMetadata() {
     return Objects.hash(keyFieldSecondary, getKeyClassSecondary());
   }
 }
diff --git a/scio-smb/src/test/java/org/apache/beam/sdk/extensions/smb/BucketMetadataTest.java b/scio-smb/src/test/java/org/apache/beam/sdk/extensions/smb/BucketMetadataTest.java
@@ -75,12 +75,12 @@ public Void extractKeySecondary(Object value) {
               }
 
               @Override
-              public int hashPrimaryKeyMetadata() {
+              int hashPrimaryKeyMetadata() {
                 return -1;
               }
 
               @Override
-              public int hashSecondaryKeyMetadata() {
+              int hashSecondaryKeyMetadata() {
                 return -1;
               }
             });

diff --git a/scio-smb/src/test/java/org/apache/beam/sdk/extensions/smb/TestBucketMetadata.java b/scio-smb/src/test/java/org/apache/beam/sdk/extensions/smb/TestBucketMetadata.java
@@ -99,12 +99,12 @@ public Void extractKeySecondary(final String value) {
   }
 
   @Override
-  public int hashPrimaryKeyMetadata() {
+  int hashPrimaryKeyMetadata() {
     return Objects.hash(getClass(), keyIndex);
   }
 
   @Override
-  public int hashSecondaryKeyMetadata() {
+  int hashSecondaryKeyMetadata() {
     throw new IllegalArgumentException();
   }
 }
diff --git a/...smb/src/test/java/org/apache/beam/sdk/extensions/smb/TestBucketMetadataWithSecondary.java b/...smb/src/test/java/org/apache/beam/sdk/extensions/smb/TestBucketMetadataWithSecondary.java
@@ -97,12 +97,12 @@ public String extractKeySecondary(final String value) {
   }
 
   @Override
-  public int hashPrimaryKeyMetadata() {
+  int hashPrimaryKeyMetadata() {
     return Objects.hash(keyIndex);
   }
 
   @Override
-  public int hashSecondaryKeyMetadata() {
+  int hashSecondaryKeyMetadata() {
     return Objects.hash(keyIndexSecondary);
   }
 }