fix: ETL job failure when input dataset is empty (#1490)

Author: amliuyong

src/data-pipeline/spark-etl/src/main/java/software/aws/solution/clickstream/TransformerV3.java

@@ -261,11 +261,6 @@ public Dataset<Row> extractUser(final Dataset<Row> eventDataset, final Dataset<R
         Dataset<Row> userDataset = convertedDataset.select(expr("dataOut.user.*"))
                 .select(toColumnArray(ModelV2.getUserFields()));
 
-        if (userDataset.count() == 0) {
-            log.info("extractUser return empty dataset");
-            return userDataset;
-        }
-
         userDataset = userDataset.withColumn(USER_FIRST_EVENT_NAME, col(Constant.EVENT_NAME))
                 .withColumn(USER_LATEST_EVENT_NAME, col(Constant.EVENT_NAME))
                 .drop(Constant.EVENT_NAME);
@@ -279,6 +274,11 @@ public Dataset<Row> extractUser(final Dataset<Row> eventDataset, final Dataset<R
         log.info("tableName: {}", tableName);
         log.info("pathInfo - incremental: " + pathInfo.getIncremental() + ", full: " + pathInfo.getFull());
 
+        if (userDataset.count() == 0) {
+            log.info("extractUser return empty dataset");
+            return userDataset;
+        }
+
         // save new (append)
         String path = saveIncrementalDatasetToPath(pathInfo.getIncremental(), newUserAggDataset);
 

src/data-pipeline/spark-etl/src/main/java/software/aws/solution/clickstream/transformer/BaseThirdPartyTransformer.java

@@ -217,11 +217,6 @@ public Dataset<Row> extractUser(final Dataset<Row> eventDataset, final Dataset<R
         Dataset<Row> userDataset = convertedDataset.select(expr("dataOut.user.*"))
                 .select(toColumnArray(ModelV2.getUserFields()));
 
-        if (userDataset.count() == 0) {
-            log.info("extractUser return empty dataset");
-            return userDataset;
-        }
-
         // agg new
         Dataset<Row> newUserAggDataset = aggUserDataset(userDataset, "newUserAggDataset");
         log.info("newUserAggDataset count: {}", newUserAggDataset.count());
@@ -231,6 +226,11 @@ public Dataset<Row> extractUser(final Dataset<Row> eventDataset, final Dataset<R
         log.info("tableName: {}", tableName);
         log.info("pathInfo - incremental: " + pathInfo.getIncremental() + ", full: " + pathInfo.getFull());
 
+        if (userDataset.count() == 0) {
+            log.info("extractUser return empty dataset");
+            return userDataset;
+        }
+
         // save new (append)
         String path = saveIncrementalDatasetToPath(pathInfo.getIncremental(), newUserAggDataset);
 

src/data-pipeline/spark-etl/src/test/java/software/aws/solution/clickstream/TransformerV3Test.java

@@ -316,4 +316,37 @@ void test_extract_session_from_event2() throws IOException {
         String expectedJson1 = this.resourceFileAsString("/event_v2/expected/test_extract_session_from_event2.json");
         Assertions.assertEquals(expectedJson1, replaceDynData(sessionDataset.first().prettyJson()));
     }
+
+
+    @Test
+    public void should_transform_can_merge_state_table_when_dataset_is_empty_v2() throws IOException {
+        // DOWNLOAD_FILE=0 ./gradlew clean test --info --tests software.aws.solution.clickstream.TransformerV3Test.should_transform_can_merge_state_table_when_dataset_is_empty_v2
+        System.setProperty(APP_IDS_PROP, "uba-app");
+        System.setProperty(PROJECT_ID_PROP, "test_project_id_01");
+        System.setProperty("force.merge", "true");
+
+        String testWarehouseDir = "/tmp/warehouse/should_transform_can_merge_state_table_when_dataset_is_empty_v2/" + new Date().getTime();
+        System.setProperty(WAREHOUSE_DIR_PROP, testWarehouseDir);
+
+        Dataset<Row> dataset =
+                spark.read().json(requireNonNull(getClass().getResource("/original_data_with_user_profile_set2.json")).getPath());
+
+        Dataset<Row> emptyDataset = dataset.filter(col("appId").equalTo("empty"));
+        // make sure input dataset is empty
+        Assertions.assertEquals(0, emptyDataset.count());
+
+
+        // create folders for full and incremental table
+        String tableDirFull = testWarehouseDir + "/" + getUserPropsTableName("clickstream") + FULL_SUFFIX  + TABLE_VERSION_SUFFIX_V3;
+        emptyDataset.write().mode(SaveMode.Overwrite).json(tableDirFull);
+        String tableDirIncr = testWarehouseDir + "/" + getUserPropsTableName("clickstream") + INCREMENTAL_SUFFIX  + TABLE_VERSION_SUFFIX_V3;
+        emptyDataset.write().mode(SaveMode.Overwrite).json(tableDirIncr);
+
+
+        Map<TableName, Dataset<Row>> transformedDatasets = transformer.transform(emptyDataset);
+        Dataset<Row> datasetEvent = transformedDatasets.get(TableName.EVENT_V2);
+
+        transformer.postTransform(datasetEvent);  // should no error
+    }
+
 }