From d7dff95481d3e3f1a2ee1256a7d7aff3cb3832ca Mon Sep 17 00:00:00 2001
From: Alec Huang <alec.huang@snowflake.com>
Date: Fri, 18 Oct 2024 12:59:50 -0700
Subject: [PATCH] Address comments 2

---
 .../streaming/internal/ParquetRowBuffer.java  | 22 +++++++---------
 .../streaming/internal/RowBufferStats.java    | 26 +++++++++----------
 2 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java b/src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java
index 58d0aa069..c2197da24 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java
@@ -358,18 +358,16 @@ private float addRow(
         if (subColumnFinder == null) {
           throw new SFException(ErrorCode.INTERNAL_ERROR, "SubColumnFinder is not initialized.");
         }
-        subColumnFinder
-            .getSubColumns(columnName)
-            .forEach(
-                subColumn -> {
-                  RowBufferStats stats = statsMap.get(subColumn);
-                  if (stats == null) {
-                    throw new SFException(
-                        ErrorCode.INTERNAL_ERROR,
-                        String.format("Column %s not found in stats map.", subColumn));
-                  }
-                  statsMap.get(subColumn).incCurrentNullCount();
-                });
+
+        for (String subColumn : subColumnFinder.getSubColumns(columnName)) {
+          RowBufferStats stats = statsMap.get(subColumn);
+          if (stats == null) {
+            throw new SFException(
+                ErrorCode.INTERNAL_ERROR,
+                String.format("Column %s not found in stats map.", subColumn));
+          }
+          stats.incCurrentNullCount();
+        }
       } else {
         statsMap.get(columnName).incCurrentNullCount();
       }
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/RowBufferStats.java b/src/main/java/net/snowflake/ingest/streaming/internal/RowBufferStats.java
index d4e986cdb..6d61c8eea 100644
--- a/src/main/java/net/snowflake/ingest/streaming/internal/RowBufferStats.java
+++ b/src/main/java/net/snowflake/ingest/streaming/internal/RowBufferStats.java
@@ -121,22 +121,22 @@ RowBufferStats forkEmpty() {
 
   // TODO performance test this vs in place update
   static RowBufferStats getCombinedStats(RowBufferStats left, RowBufferStats right) {
-    if (!Objects.equals(left.getCollationDefinitionString(), right.collationDefinitionString)) {
+    if (!Objects.equals(left.getCollationDefinitionString(), right.collationDefinitionString)
+        || left.enableDistinctValuesCount != right.enableDistinctValuesCount
+        || left.enableValuesCount != right.enableValuesCount) {
       throw new SFException(
           ErrorCode.INVALID_COLLATION_STRING,
-          "Tried to combine stats for different collations",
+          "Tried to combine stats for different"
+              + " collations/enableDistinctValuesCount/enableValuesCount",
           String.format(
-              "left=%s, right=%s",
-              left.getCollationDefinitionString(), right.getCollationDefinitionString()));
-    }
-
-    if (left.enableDistinctValuesCount != right.enableDistinctValuesCount) {
-      throw new SFException(
-          ErrorCode.INTERNAL_ERROR,
-          "Tried to combine stats for different distinct value settings",
-          String.format(
-              "left=%s, right=%s",
-              left.enableDistinctValuesCount, right.enableDistinctValuesCount));
+              "left={collations=%s, enableDistinctValuesCount=%s, enableValuesCount=%s}, "
+                  + "right={collations=%s, enableDistinctValuesCount=%s, enableValuesCount=%s}",
+              left.getCollationDefinitionString(),
+              left.enableDistinctValuesCount,
+              left.enableValuesCount,
+              right.getCollationDefinitionString(),
+              right.enableDistinctValuesCount,
+              right.enableValuesCount));
     }
 
     RowBufferStats combined =