SNOW-928922 send spansMixedTables flag in blob registration requests (#…

…651)
snowflakedb · Dec 18, 2023 · e140aa7 · e140aa7
1 parent 2b702db
commit e140aa7
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 7 deletions.
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/BlobMetadata.java b/src/main/java/net/snowflake/ingest/streaming/internal/BlobMetadata.java
@@ -18,24 +18,33 @@ class BlobMetadata {
   private final Constants.BdecVersion bdecVersion;
   private final List<ChunkMetadata> chunks;
   private final BlobStats blobStats;
+  private final boolean spansMixedTables;
 
   // used for testing only
   @VisibleForTesting
   BlobMetadata(String path, String md5, List<ChunkMetadata> chunks, BlobStats blobStats) {
-    this(path, md5, ParameterProvider.BLOB_FORMAT_VERSION_DEFAULT, chunks, blobStats);
+    this(
+        path,
+        md5,
+        ParameterProvider.BLOB_FORMAT_VERSION_DEFAULT,
+        chunks,
+        blobStats,
+        chunks == null ? false : chunks.size() > 1);
   }
 
   BlobMetadata(
       String path,
       String md5,
       Constants.BdecVersion bdecVersion,
       List<ChunkMetadata> chunks,
-      BlobStats blobStats) {
+      BlobStats blobStats,
+      boolean spansMixedTables) {
     this.path = path;
     this.md5 = md5;
     this.bdecVersion = bdecVersion;
     this.chunks = chunks;
     this.blobStats = blobStats;
+    this.spansMixedTables = spansMixedTables;
   }
 
   @JsonIgnore
@@ -68,13 +77,19 @@ BlobStats getBlobStats() {
     return this.blobStats;
   }
 
+  @JsonProperty("spans_mixed_tables")
+  boolean getSpansMixedTables() {
+    return this.spansMixedTables;
+  }
+
   /** Create {@link BlobMetadata}. */
   static BlobMetadata createBlobMetadata(
       String path,
       String md5,
       Constants.BdecVersion bdecVersion,
       List<ChunkMetadata> chunks,
-      BlobStats blobStats) {
-    return new BlobMetadata(path, md5, bdecVersion, chunks, blobStats);
+      BlobStats blobStats,
+      boolean spansMixedTables) {
+    return new BlobMetadata(path, md5, bdecVersion, chunks, blobStats, spansMixedTables);
   }
 }
diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java b/src/main/java/net/snowflake/ingest/streaming/internal/FlushService.java
@@ -568,8 +568,15 @@ BlobMetadata upload(
         blob.length,
         System.currentTimeMillis() - startTime);
 
+    // at this point we know for sure if the BDEC file has data for more than one chunk, i.e.
+    // spans mixed tables or not
     return BlobMetadata.createBlobMetadata(
-        blobPath, BlobBuilder.computeMD5(blob), bdecVersion, metadata, blobStats);
+        blobPath,
+        BlobBuilder.computeMD5(blob),
+        bdecVersion,
+        metadata,
+        blobStats,
+        metadata == null ? false : metadata.size() > 1);
   }
 
   /**

diff --git a/.../java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java b/.../java/net/snowflake/ingest/streaming/internal/SnowflakeStreamingIngestClientInternal.java
@@ -719,7 +719,10 @@ List<BlobMetadata> getRetryBlobs(
                     blobMetadata.getMD5(),
                     blobMetadata.getVersion(),
                     relevantChunks,
-                    blobMetadata.getBlobStats()));
+                    blobMetadata.getBlobStats(),
+                    // Important to not change the spansMixedTables value in case of retries. The
+                    // correct value is the value that the already uploaded blob has.
+                    blobMetadata.getSpansMixedTables()));
           }
         });
 

diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/RegisterServiceTest.java
@@ -7,6 +7,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import net.snowflake.ingest.utils.Pair;
@@ -17,7 +18,7 @@
 public class RegisterServiceTest {
 
   @Test
-  public void testRegisterService() {
+  public void testRegisterService() throws ExecutionException, InterruptedException {
     RegisterService<StubChunkData> rs = new RegisterService<>(null, true);
 
     Pair<FlushService.BlobData<StubChunkData>, CompletableFuture<BlobMetadata>> blobFuture =
@@ -26,6 +27,7 @@ public void testRegisterService() {
             CompletableFuture.completedFuture(new BlobMetadata("path", "md5", null, null)));
     rs.addBlobs(Collections.singletonList(blobFuture));
     Assert.assertEquals(1, rs.getBlobsList().size());
+    Assert.assertEquals(false, blobFuture.getValue().get().getSpansMixedTables());
     List<FlushService.BlobData<StubChunkData>> errorBlobs = rs.registerBlobs(null);
     Assert.assertEquals(0, rs.getBlobsList().size());
     Assert.assertEquals(0, errorBlobs.size());