[s] update sample size to int

cloudsufi · Jan 8, 2024 · 2f35ce9 · 2f35ce9
1 parent e002f14
commit 2f35ce9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 11 deletions.
diff --git a/format-xls/src/main/java/io/cdap/plugin/format/xls/input/XlsInputFormatConfig.java b/format-xls/src/main/java/io/cdap/plugin/format/xls/input/XlsInputFormatConfig.java
@@ -107,8 +107,13 @@ public XlsInputFormatConfig(@Nullable String schema, @Nullable String sheet, @Nu
     this.terminateIfEmptyRow = terminateIfEmptyRow;
   }
 
-  public long getSampleSize() {
-    return Long.parseLong(getProperties().getProperties().getOrDefault(NAME_SAMPLE_SIZE, "1000"));
+  public int getSampleSize() {
+    String sampleSize = getProperties().getProperties().getOrDefault(NAME_SAMPLE_SIZE, "1000");
+    try {
+      return Integer.parseInt(sampleSize);
+    } catch (NumberFormatException e) {
+      throw new IllegalArgumentException(String.format("Invalid sample size '%s'.", sampleSize));
+    }
   }
 
   public String getSheet() {

diff --git a/format-xls/src/main/java/io/cdap/plugin/format/xls/input/XlsInputFormatProvider.java b/format-xls/src/main/java/io/cdap/plugin/format/xls/input/XlsInputFormatProvider.java
@@ -127,10 +127,10 @@ public Schema detectSchema(FormatContext context, InputFiles inputFiles) throws
           return null;
         }
 
-        int sampleSizeInt = getSampleSizeInt();
+        int sampleSize = conf.getSampleSize();
         // Row numbers are 0 based in POI
         int rowStart = Math.min(0, workSheet.getFirstRowNum());
-        int rowEnd = Math.min(sampleSizeInt, workSheet.getLastRowNum());
+        int rowEnd = Math.min(sampleSize, workSheet.getLastRowNum());
 
         int lastCellNumMax = 0;
         List<String> columnNames = new ArrayList<>();
@@ -182,11 +182,4 @@ public Schema detectSchema(FormatContext context, InputFiles inputFiles) throws
     return null;
   }
 
-  private int getSampleSizeInt() {
-    // Note: Sample size is long, but we are casting it to int here. This is because the POI API uses int
-    // for row numbers, so we cannot support sample sizes larger than Integer.MAX_VALUE
-    long sampleSize = conf.getSampleSize() > Integer.MAX_VALUE ? Integer.MAX_VALUE : conf.getSampleSize();
-    return (int) sampleSize;
-  }
-
 }