SNOW-1480787 Fix DataFrameSuite test to validate with pivot v1 and pi…

…vot v2
snowflakedb · Jul 29, 2024 · 302895d · 302895d
1 parent bfbddd6
commit 302895d
Showing 1 changed file with 27 additions and 11 deletions.
diff --git a/src/test/scala/com/snowflake/snowpark_test/DataFrameSuite.scala b/src/test/scala/com/snowflake/snowpark_test/DataFrameSuite.scala
@@ -469,9 +469,9 @@ trait DataFrameSuite extends TestData with BeforeAndAfterEach {
     assert(double2.stat.approxQuantile(Array[String](), Array[Double]()).isEmpty)
   }
 
-  test("df.stat.pivot") {
+  private def testDataframeStatPivot() : Unit = {
     assert(
-      getShowString(monthlySales.stat.crosstab("empid", "month"), 10) ==
+      getShowString(monthlySales.stat.crosstab("empid", "month").sort(col("empid")), 10) ==
         """---------------------------------------------------
           ||"EMPID"  |"'JAN'"  |"'FEB'"  |"'MAR'"  |"'APR'"  |
           |---------------------------------------------------
@@ -481,29 +481,29 @@ trait DataFrameSuite extends TestData with BeforeAndAfterEach {
           |""".stripMargin)
 
     assert(
-      getShowString(monthlySales.stat.crosstab("month", "empid"), 10) ==
+      getShowString(monthlySales.stat.crosstab("month", "empid").sort(col("month")), 10) ==
         """-------------------------------------------------------------------
           ||"MONTH"  |"CAST(1 AS NUMBER(38,0))"  |"CAST(2 AS NUMBER(38,0))"  |
           |-------------------------------------------------------------------
-          ||JAN      |2                          |2                          |
+          ||APR      |2                          |2                          |
           ||FEB      |2                          |2                          |
+          ||JAN      |2                          |2                          |
           ||MAR      |2                          |2                          |
-          ||APR      |2                          |2                          |
           |-------------------------------------------------------------------
           |""".stripMargin)
 
     assert(
-      getShowString(date1.stat.crosstab("a", "b"), 10) ==
+      getShowString(date1.stat.crosstab("a", "b").sort(col("a")), 10) ==
         """----------------------------------------------------------------------
           ||"A"         |"CAST(1 AS NUMBER(38,0))"  |"CAST(2 AS NUMBER(38,0))"  |
           |----------------------------------------------------------------------
-          ||2020-08-01  |1                          |0                          |
           ||2010-12-01  |0                          |1                          |
+          ||2020-08-01  |1                          |0                          |
           |----------------------------------------------------------------------
           |""".stripMargin)
 
     assert(
-      getShowString(date1.stat.crosstab("b", "a"), 10) ==
+      getShowString(date1.stat.crosstab("b", "a").sort(col("b")), 10) ==
         """-----------------------------------------------------------
           ||"B"  |"TO_DATE('2020-08-01')"  |"TO_DATE('2010-12-01')"  |
           |-----------------------------------------------------------
@@ -513,17 +513,17 @@ trait DataFrameSuite extends TestData with BeforeAndAfterEach {
           |""".stripMargin)
 
     assert(
-      getShowString(string7.stat.crosstab("a", "b"), 10) ==
+      getShowString(string7.stat.crosstab("a", "b").sort(col("a")), 10) ==
         """----------------------------------------------------------------
           ||"A"   |"CAST(1 AS NUMBER(38,0))"  |"CAST(2 AS NUMBER(38,0))"  |
           |----------------------------------------------------------------
-          ||str   |1                          |0                          |
           ||NULL  |0                          |1                          |
+          ||str   |1                          |0                          |
           |----------------------------------------------------------------
           |""".stripMargin)
 
     assert(
-      getShowString(string7.stat.crosstab("b", "a"), 10) ==
+      getShowString(string7.stat.crosstab("b", "a").sort(col("b")), 10) ==
         """--------------------------
           ||"B"  |"'str'"  |"NULL"  |
           |--------------------------
@@ -533,6 +533,22 @@ trait DataFrameSuite extends TestData with BeforeAndAfterEach {
           |""".stripMargin)
   }
 
+  test("df.stat.pivot") {
+    testWithAlteredSessionParameter(
+      testDataframeStatPivot(),
+      "ENABLE_PIVOT_VIEW_WITH_OBJECT_AGG",
+      "disable",
+      skipIfParamNotExist = true)
+  }
+
+  test("df.stat.pivot.2") {
+    testWithAlteredSessionParameter(
+      testDataframeStatPivot(),
+      "ENABLE_PIVOT_VIEW_WITH_OBJECT_AGG",
+      "enable",
+      skipIfParamNotExist = true)
+  }
+
   test("df.stat.sampleBy") {
     assert(
       getShowString(monthlySales.stat.sampleBy(col("empid"), Map(1 -> 0.0, 2 -> 1.0)), 10) ==