modifying test cases

data-integrations · Nov 14, 2023 · f11cb6b · f11cb6b
1 parent b2c0b9c
commit f11cb6b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 13 deletions.
diff --git a/wrangler-core/src/main/java/io/cdap/directives/column/SetType.java b/wrangler-core/src/main/java/io/cdap/directives/column/SetType.java
@@ -177,12 +177,12 @@ public Schema getOutputSchema(SchemaResolutionContext context) {
   /**
    * extracts precision and scale from schema string
    */
-    public static Pair<Integer, Integer> findPrecisionAndScale(Schema inputSchema) {
+    public static Pair<Integer, Integer> findPrecisionAndScale(Schema intermediateSchema) {
       Integer precision = null;
       Integer scale = null;
-      if (inputSchema.getLogicalType() == LogicalType.DECIMAL) {
-      precision = inputSchema.getPrecision();
-      scale = inputSchema.getScale();
+      if (intermediateSchema.getLogicalType() == LogicalType.DECIMAL) {
+      precision = intermediateSchema.getPrecision();
+      scale = intermediateSchema.getScale();
       }
       return new Pair<Integer, Integer>(precision, scale);
     }

diff --git a/wrangler-core/src/test/java/io/cdap/directives/column/SetTypeTest.java b/wrangler-core/src/test/java/io/cdap/directives/column/SetTypeTest.java
@@ -18,6 +18,7 @@
 
 import io.cdap.cdap.api.common.Bytes;
 import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.api.data.schema.Schema.Type;
 import io.cdap.wrangler.TestingRig;
 import io.cdap.wrangler.api.RecipeException;
 import io.cdap.wrangler.api.Row;
@@ -27,6 +28,7 @@
 import java.math.BigDecimal;
 import java.time.LocalDateTime;
 import java.time.ZonedDateTime;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
@@ -232,7 +234,7 @@ public void testToDecimalWithScalePrecisionAndRoundingMode() throws Exception {
   @Test
   public void testToDecimalWithPrecision() throws Exception {
     List<Row> rows = Collections.singletonList(new Row("scale_1_precision_4", "122.5"));
-    String[] directives = new String[] {"set-type :scale_1_precision_4 decimal 0 'FLOOR' prop:{precision=3}"};
+    String[] directives = new String[] {"set-type :scale_1_precision_4 decimal 'FLOOR' prop:{precision=3}"};
     List<Row> results = TestingRig.execute(directives, rows);
     Row row = results.get(0);
 
@@ -246,14 +248,18 @@ public void testToDecimalWithInvalidPrecision() throws Exception {
     List<Row> rows = Collections.singletonList(new Row("scale_1_precision_4", "122.5"));
     String[] directives = new String[] {"set-type :scale_1_precision_4 decimal 0 'FLOOR' prop:{precision=-1}"};
     TestingRig.execute(directives, rows);
+  }
 
   @Test
   public void testToDecimalScaleIsNull() throws Exception {
     List<Row> rows = Collections.singletonList(new Row("scale_2", "125.45"));
     String[] directives = new String[] {"set-type scale_2 decimal"};
+    List<Schema> unionSchemaList = new ArrayList<>();
+    unionSchemaList.add(Schema.of(Type.DOUBLE));
+    unionSchemaList.add(Schema.of(Type.NULL));
     Schema inputSchema = Schema.recordOf(
         "inputSchema",
-        Schema.Field.of("scale_2", Schema.of(Schema.Type.DOUBLE))
+        Schema.Field.of("scale_2", Schema.unionOf(unionSchemaList))
     );
 
     Schema expectedSchema = Schema.recordOf(
@@ -412,13 +418,13 @@ public void testGetOutputSchemaForTypeChangedColumn() throws Exception {
     );
     Schema inputSchema = Schema.recordOf(
       "inputSchema",
-      Schema.Field.of("A", Schema.of(Schema.Type.STRING)),
-      Schema.Field.of("B", Schema.of(Schema.Type.STRING)),
-      Schema.Field.of("C", Schema.of(Schema.Type.STRING)),
-      Schema.Field.of("D", Schema.of(Schema.Type.STRING)),
-      Schema.Field.of("E", Schema.of(Schema.Type.INT)),
-      Schema.Field.of("F", Schema.of(Schema.Type.STRING)),
-      Schema.Field.of("G", Schema.of(Schema.Type.LONG))
+      Schema.Field.of("A", Schema.unionOf(Schema.of(Schema.Type.STRING), Schema.of(Type.NULL))),
+      Schema.Field.of("B", Schema.unionOf(Schema.of(Schema.Type.STRING), Schema.of(Type.NULL))),
+      Schema.Field.of("C", Schema.unionOf(Schema.of(Schema.Type.STRING), Schema.of(Type.NULL))),
+      Schema.Field.of("D", Schema.unionOf(Schema.of(Schema.Type.STRING), Schema.of(Type.NULL))),
+      Schema.Field.of("E", Schema.unionOf(Schema.of(Schema.Type.STRING), Schema.of(Type.NULL))),
+      Schema.Field.of("F", Schema.unionOf(Schema.of(Schema.Type.STRING), Schema.of(Type.NULL))),
+      Schema.Field.of("G", Schema.unionOf(Schema.of(Schema.Type.STRING), Schema.of(Type.NULL)))
     );
     Schema expectedSchema = Schema.recordOf(
       "expectedSchema",