Add tests with numThreadsForFinalReduce

apache · Jan 18, 2025 · fcb643d · fcb643d
1 parent 2fba98b
commit fcb643d
Show file tree

Hide file tree

Showing 18 changed files with 107 additions and 46 deletions.
diff --git a/...-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ArrayTest.java b/...-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ArrayTest.java
@@ -873,7 +873,7 @@ public Schema createSchema() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // create avro schema
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
@@ -953,6 +953,6 @@ public File createAvroFile()
         ));
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 }
diff --git a/...egration-tests/src/test/java/org/apache/pinot/integration/tests/custom/BytesTypeTest.java b/...egration-tests/src/test/java/org/apache/pinot/integration/tests/custom/BytesTypeTest.java
@@ -23,6 +23,7 @@
 import java.io.File;
 import java.nio.ByteBuffer;
 import java.util.Base64;
+import java.util.List;
 import java.util.UUID;
 import org.apache.avro.file.DataFileWriter;
 import org.apache.avro.generic.GenericData;
@@ -88,7 +89,7 @@ protected long getCountStarResult() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // create avro schema
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
@@ -153,8 +154,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-
-    return avroFile;
+    return List.of(avroFile);
   }
 
   private static String newRandomBase64String() {

diff --git a/...egration-tests/src/test/java/org/apache/pinot/integration/tests/custom/CpcSketchTest.java b/...egration-tests/src/test/java/org/apache/pinot/integration/tests/custom/CpcSketchTest.java
@@ -23,6 +23,7 @@
 import java.io.File;
 import java.nio.ByteBuffer;
 import java.util.Base64;
+import java.util.List;
 import java.util.Random;
 import org.apache.avro.file.DataFileWriter;
 import org.apache.avro.generic.GenericData;
@@ -165,7 +166,7 @@ public Schema createSchema() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // create avro schema
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
@@ -188,7 +189,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   private byte[] getRandomRawValue() {

diff --git a/...java/org/apache/pinot/integration/tests/custom/CustomDataQueryClusterIntegrationTest.java b/...java/org/apache/pinot/integration/tests/custom/CustomDataQueryClusterIntegrationTest.java
@@ -94,22 +94,26 @@ public void setUp()
     Schema schema = createSchema();
     addSchema(schema);
 
-    File avroFile = createAvroFile();
+    List<File> avroFiles = createAvroFiles();
     if (isRealtimeTable()) {
       // create realtime table
-      TableConfig tableConfig = createRealtimeTableConfig(avroFile);
+      TableConfig tableConfig = createRealtimeTableConfig(avroFiles.get(0));
       addTableConfig(tableConfig);
 
       // Push data into Kafka
-      pushAvroIntoKafka(List.of(avroFile));
+      pushAvroIntoKafka(avroFiles);
     } else {
       // create offline table
       TableConfig tableConfig = createOfflineTableConfig();
       addTableConfig(tableConfig);
 
       // create & upload segments
-      ClusterIntegrationTestUtils.buildSegmentFromAvro(avroFile, tableConfig, schema, 0, _segmentDir, _tarDir);
-      uploadSegments(getTableName(), _tarDir);
+      int segmentIndex = 0;
+      for (File avroFile : avroFiles) {
+        ClusterIntegrationTestUtils.buildSegmentFromAvro(avroFile, tableConfig, schema, segmentIndex++, _segmentDir,
+            _tarDir);
+        uploadSegments(getTableName(), _tarDir);
+      }
     }
 
     waitForAllDocsLoaded(60_000);
@@ -247,7 +251,7 @@ public String getKafkaTopic() {
   @Override
   public abstract Schema createSchema();
 
-  public abstract File createAvroFile()
+  public abstract List<File> createAvroFiles()
       throws Exception;
 
   public boolean isRealtimeTable() {

diff --git a/...ts/src/test/java/org/apache/pinot/integration/tests/custom/FloatingPointDataTypeTest.java b/...ts/src/test/java/org/apache/pinot/integration/tests/custom/FloatingPointDataTypeTest.java
@@ -72,7 +72,7 @@ public Schema createSchema() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws IOException {
 
     // create avro schema
@@ -124,7 +124,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   @Override

diff --git a/...gration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GeoSpatialTest.java b/...gration-tests/src/test/java/org/apache/pinot/integration/tests/custom/GeoSpatialTest.java
@@ -22,6 +22,7 @@
 import com.google.common.collect.ImmutableList;
 import java.io.File;
 import java.nio.ByteBuffer;
+import java.util.List;
 import org.apache.avro.file.DataFileWriter;
 import org.apache.avro.generic.GenericData;
 import org.apache.avro.generic.GenericDatumWriter;
@@ -130,7 +131,7 @@ protected long getCountStarResult() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // create avro schema
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
@@ -185,7 +186,7 @@ public File createAvroFile()
       }
     }
 
-    return avroFile;
+    return List.of(avroFile);
   }
 
   @Test(dataProvider = "useBothQueryEngines")

diff --git a/...tegration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonPathTest.java b/...tegration-tests/src/test/java/org/apache/pinot/integration/tests/custom/JsonPathTest.java
@@ -96,7 +96,7 @@ public TableConfig createOfflineTableConfig() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
     List<org.apache.avro.Schema.Field> fields =
@@ -130,7 +130,7 @@ public File createAvroFile()
     }
     Collections.sort(_sortedSequenceIds);
 
-    return avroFile;
+    return List.of(avroFile);
   }
 
   @Test(dataProvider = "useBothQueryEngines")

diff --git a/...sts/src/test/java/org/apache/pinot/integration/tests/custom/MapFieldTypeRealtimeTest.java b/...sts/src/test/java/org/apache/pinot/integration/tests/custom/MapFieldTypeRealtimeTest.java
@@ -92,7 +92,7 @@ public TableConfig createOfflineTableConfig() {
         .build();
   }
 
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
     org.apache.avro.Schema stringMapAvroSchema =
@@ -126,7 +126,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   protected int getSelectionDefaultDocCount() {

diff --git a/...ation-tests/src/test/java/org/apache/pinot/integration/tests/custom/MapFieldTypeTest.java b/...ation-tests/src/test/java/org/apache/pinot/integration/tests/custom/MapFieldTypeTest.java
@@ -89,7 +89,7 @@ public TableConfig createOfflineTableConfig() {
         .build();
   }
 
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
     org.apache.avro.Schema stringMapAvroSchema =
@@ -119,7 +119,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   protected int getSelectionDefaultDocCount() {

diff --git a/...ntegration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MapTypeTest.java b/...ntegration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MapTypeTest.java
@@ -86,7 +86,7 @@ public TableConfig createOfflineTableConfig() {
         .build();
   }
 
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
     org.apache.avro.Schema stringKeyMapAvroSchema =
@@ -116,7 +116,7 @@ public File createAvroFile()
       }
     }
 
-    return avroFile;
+    return List.of(avroFile);
   }
 
   protected int getSelectionDefaultDocCount() {

diff --git a/...ation-tests/src/test/java/org/apache/pinot/integration/tests/custom/SumPrecisionTest.java b/...ation-tests/src/test/java/org/apache/pinot/integration/tests/custom/SumPrecisionTest.java
@@ -24,6 +24,7 @@
 import java.io.IOException;
 import java.math.BigDecimal;
 import java.nio.ByteBuffer;
+import java.util.List;
 import java.util.Random;
 import org.apache.avro.file.DataFileWriter;
 import org.apache.avro.generic.GenericData;
@@ -64,7 +65,7 @@ public Schema createSchema() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws IOException {
 
     // create avro schema
@@ -103,7 +104,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   @Override

diff --git a/...ration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TextIndicesTest.java b/...ration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TextIndicesTest.java
@@ -132,7 +132,7 @@ protected long getCountStarResult() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // Read all skills from the skill file
     InputStream inputStream = getClass().getClassLoader().getResourceAsStream("data/text_search_data/skills.txt");
@@ -164,7 +164,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   @Override

diff --git a/...ration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ThetaSketchTest.java b/...ration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ThetaSketchTest.java
@@ -87,7 +87,7 @@ protected long getCountStarResult() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws IOException {
 
     // create avro schema
@@ -171,7 +171,7 @@ public File createAvroFile()
       }
     }
 
-    return avroFile;
+    return List.of(avroFile);
   }
 
   @Test(dataProvider = "useV1QueryEngine")

diff --git a/...egration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TimestampTest.java b/...egration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TimestampTest.java
@@ -22,6 +22,7 @@
 import com.google.common.collect.ImmutableList;
 import java.io.File;
 import java.sql.Timestamp;
+import java.util.List;
 import java.util.TimeZone;
 import org.apache.avro.Schema.Field;
 import org.apache.avro.Schema.Type;
@@ -461,7 +462,7 @@ public Schema createSchema() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // create avro schema
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
@@ -533,6 +534,6 @@ public File createAvroFile()
         tsBaseLong += 86400000;
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 }
diff --git a/...ration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TupleSketchTest.java b/...ration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TupleSketchTest.java
@@ -23,6 +23,7 @@
 import java.io.File;
 import java.nio.ByteBuffer;
 import java.util.Base64;
+import java.util.List;
 import java.util.Random;
 import org.apache.avro.file.DataFileWriter;
 import org.apache.avro.generic.GenericData;
@@ -279,7 +280,7 @@ public Schema createSchema() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // create avro schema
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
@@ -303,7 +304,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   private byte[] getRandomRawValue() {

diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ULLTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/ULLTest.java
@@ -24,6 +24,7 @@
 import java.io.File;
 import java.nio.ByteBuffer;
 import java.util.Base64;
+import java.util.List;
 import java.util.Random;
 import org.apache.avro.file.DataFileWriter;
 import org.apache.avro.generic.GenericData;
@@ -123,7 +124,7 @@ public Schema createSchema() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // create avro schema
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
@@ -146,7 +147,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   private byte[] getRandomRawValue() {

diff --git a/...integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VectorTest.java b/...integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/VectorTest.java
@@ -252,7 +252,7 @@ public Schema createSchema() {
   }
 
   @Override
-  public File createAvroFile()
+  public List<File> createAvroFiles()
       throws Exception {
     // create avro schema
     org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
@@ -320,7 +320,7 @@ public File createAvroFile()
         fileWriter.append(record);
       }
     }
-    return avroFile;
+    return List.of(avroFile);
   }
 
   private float[] createZeroVector(int vectorDimSize) {