Address comments

snowflakedb · Oct 8, 2024 · fbc4210 · fbc4210
1 parent 510f720
commit fbc4210
Show file tree

Hide file tree

Showing 11 changed files with 829 additions and 728 deletions.
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/RowBufferStats.java b/src/main/java/net/snowflake/ingest/streaming/internal/RowBufferStats.java
@@ -60,6 +60,10 @@ class RowBufferStats {
     this(columnDisplayName, null, -1, null, null);
   }
 
+  RowBufferStats(String columnDisplayName, PrimitiveType primitiveType) {
+    this(columnDisplayName, null, -1, null, primitiveType);
+  }
+
   void reset() {
     this.currentMaxStrValue = null;
     this.currentMinStrValue = null;

diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/FlushServiceTest.java
@@ -52,6 +52,8 @@
 import net.snowflake.ingest.utils.ErrorCode;
 import net.snowflake.ingest.utils.ParameterProvider;
 import net.snowflake.ingest.utils.SFException;
+import org.apache.parquet.schema.PrimitiveType;
+import org.apache.parquet.schema.Types;
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -877,8 +879,12 @@ public void testBuildAndUpload() throws Exception {
     Map<String, RowBufferStats> eps1 = new HashMap<>();
     Map<String, RowBufferStats> eps2 = new HashMap<>();
 
-    RowBufferStats stats1 = new RowBufferStats("COL1");
-    RowBufferStats stats2 = new RowBufferStats("COL1");
+    RowBufferStats stats1 =
+        new RowBufferStats(
+            "COL1", Types.optional(PrimitiveType.PrimitiveTypeName.INT32).id(1).named("COL1"));
+    RowBufferStats stats2 =
+        new RowBufferStats(
+            "COL1", Types.optional(PrimitiveType.PrimitiveTypeName.INT32).id(1).named("COL1"));
 
     eps1.put("one", stats1);
     eps2.put("one", stats2);
@@ -1115,7 +1121,9 @@ public void testBlobBuilder() throws Exception {
 
     Map<String, RowBufferStats> eps1 = new HashMap<>();
 
-    RowBufferStats stats1 = new RowBufferStats("COL1");
+    RowBufferStats stats1 =
+        new RowBufferStats(
+            "COL1", Types.optional(PrimitiveType.PrimitiveTypeName.INT32).id(1).named("COL1"));
 
     eps1.put("one", stats1);
 

diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/RowBufferTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/RowBufferTest.java
@@ -29,6 +29,8 @@
 import org.apache.commons.codec.binary.Hex;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.parquet.hadoop.BdecParquetReader;
+import org.apache.parquet.schema.PrimitiveType;
+import org.apache.parquet.schema.Types;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -563,12 +565,18 @@ private void testDoubleQuotesColumnNameHelper(OpenChannelRequest.OnErrorOption o
   public void testBuildEpInfoFromStats() {
     Map<String, RowBufferStats> colStats = new HashMap<>();
 
-    RowBufferStats stats1 = new RowBufferStats("intColumn");
+    RowBufferStats stats1 =
+        new RowBufferStats(
+            "intColumn",
+            Types.optional(PrimitiveType.PrimitiveTypeName.INT32).id(1).named("intColumn"));
     stats1.addIntValue(BigInteger.valueOf(2));
     stats1.addIntValue(BigInteger.valueOf(10));
     stats1.addIntValue(BigInteger.valueOf(1));
 
-    RowBufferStats stats2 = new RowBufferStats("strColumn");
+    RowBufferStats stats2 =
+        new RowBufferStats(
+            "strColumn",
+            Types.optional(PrimitiveType.PrimitiveTypeName.BINARY).id(2).named("strColumn"));
     stats2.addStrValue("alice");
     stats2.addStrValue("bob");
     stats2.incCurrentNullCount();
@@ -603,8 +611,14 @@ public void testBuildEpInfoFromNullColumnStats() {
     final String realColName = "realCol";
     Map<String, RowBufferStats> colStats = new HashMap<>();
 
-    RowBufferStats stats1 = new RowBufferStats(intColName);
-    RowBufferStats stats2 = new RowBufferStats(realColName);
+    RowBufferStats stats1 =
+        new RowBufferStats(
+            intColName,
+            Types.optional(PrimitiveType.PrimitiveTypeName.INT32).id(1).named(intColName));
+    RowBufferStats stats2 =
+        new RowBufferStats(
+            realColName,
+            Types.optional(PrimitiveType.PrimitiveTypeName.DOUBLE).id(2).named(realColName));
     stats1.incCurrentNullCount();
     stats2.incCurrentNullCount();
 
@@ -638,12 +652,18 @@ public void testBuildEpInfoFromNullColumnStats() {
   public void testInvalidEPInfo() {
     Map<String, RowBufferStats> colStats = new HashMap<>();
 
-    RowBufferStats stats1 = new RowBufferStats("intColumn");
+    RowBufferStats stats1 =
+        new RowBufferStats(
+            "intColumn",
+            Types.optional(PrimitiveType.PrimitiveTypeName.INT32).id(1).named("intColumn"));
     stats1.addIntValue(BigInteger.valueOf(2));
     stats1.addIntValue(BigInteger.valueOf(10));
     stats1.addIntValue(BigInteger.valueOf(1));
 
-    RowBufferStats stats2 = new RowBufferStats("strColumn");
+    RowBufferStats stats2 =
+        new RowBufferStats(
+            "strColumn",
+            Types.optional(PrimitiveType.PrimitiveTypeName.BINARY).id(2).named("strColumn"));
     stats2.addStrValue("alice");
     stats2.incCurrentNullCount();
     stats2.incCurrentNullCount();

diff --git a/...test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java b/...test/java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientTest.java
@@ -61,6 +61,8 @@
 import net.snowflake.ingest.utils.SnowflakeURL;
 import net.snowflake.ingest.utils.Utils;
 import org.apache.commons.lang3.tuple.Pair;
+import org.apache.parquet.schema.PrimitiveType;
+import org.apache.parquet.schema.Types;
 import org.bouncycastle.asn1.nist.NISTObjectIdentifiers;
 import org.bouncycastle.openssl.jcajce.JcaPEMWriter;
 import org.bouncycastle.operator.OperatorCreationException;
@@ -500,7 +502,10 @@ public void testRegisterBlobRequestCreationSuccess() throws Exception {
             .build();
 
     Map<String, RowBufferStats> columnEps = new HashMap<>();
-    columnEps.put("column", new RowBufferStats("COL1"));
+    columnEps.put(
+        "column",
+        new RowBufferStats(
+            "COL1", Types.optional(PrimitiveType.PrimitiveTypeName.INT32).id(1).named("COL1")));
     EpInfo epInfo = AbstractRowBuffer.buildEpInfoFromStats(1, columnEps, !isIcebergMode);
 
     ChunkMetadata chunkMetadata =
@@ -549,7 +554,10 @@ public void testRegisterBlobRequestCreationSuccess() throws Exception {
 
   private Pair<List<BlobMetadata>, Set<ChunkRegisterStatus>> getRetryBlobMetadata() {
     Map<String, RowBufferStats> columnEps = new HashMap<>();
-    columnEps.put("column", new RowBufferStats("COL1"));
+    columnEps.put(
+        "column",
+        new RowBufferStats(
+            "COL1", Types.optional(PrimitiveType.PrimitiveTypeName.INT32).id(1).named("COL1")));
     EpInfo epInfo = AbstractRowBuffer.buildEpInfoFromStats(1, columnEps, !isIcebergMode);
 
     ChannelMetadata channelMetadata1 =

diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/datatypes/AbstractDataTypeTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/datatypes/AbstractDataTypeTest.java
@@ -62,7 +62,7 @@ public abstract class AbstractDataTypeTest {
 
   private String schemaName = "PUBLIC";
   private SnowflakeStreamingIngestClient client;
-  private static final ObjectMapper objectMapper = new ObjectMapper();
+  protected static final ObjectMapper objectMapper = new ObjectMapper();
 
   @Parameters(name = "{index}: {0}")
   public static Object[] compressionAlgorithms() {