Add com.snowflake.snowpark.functions.countDistinct (#102)

Add com.snowflake.snowpark.functions.countDistinct

src/main/java/com/snowflake/snowpark_java/Functions.java

@@ -185,6 +185,37 @@ public static Column count(Column col) {
     return new Column(com.snowflake.snowpark.functions.count(col.toScalaColumn()));
   }
 
+  /**
+   * Returns either the number of non-NULL distinct records for the specified columns, or the total
+   * number of the distinct records. An alias of count_distinct.
+   *
+   * @since 1.13.0
+   * @param first The first column name
+   * @param remaining A column name list except the first column name
+   * @return The result column
+   */
+  public static Column countDistinct(String first, String... remaining) {
+    return new Column(
+        com.snowflake.snowpark.functions.countDistinct(
+            first, JavaUtils.stringArrayToStringSeq(remaining)));
+  }
+
+  /**
+   * Returns either the number of non-NULL distinct records for the specified columns, or the total
+   * number of the distinct records. An alias of count_distinct.
+   *
+   * @since 1.13.0
+   * @param first The first column
+   * @param remaining A column list except the first column
+   * @return The result column
+   */
+  public static Column countDistinct(Column first, Column... remaining) {
+    return new Column(
+        com.snowflake.snowpark.functions.countDistinct(
+            first.toScalaColumn(),
+            JavaUtils.columnArrayToSeq(Column.toScalaColumnArray(remaining))));
+  }
+
   /**
    * Returns either the number of non-NULL distinct records for the specified columns, or the total
    * number of the distinct records.

src/main/scala/com/snowflake/snowpark/functions.scala

@@ -210,6 +210,26 @@ object functions {
     case _ => builtin("count")(e)
   }
 
+  /**
+   * Returns either the number of non-NULL distinct records for the specified columns,
+   * or the total number of the distinct records. An alias of count_distinct.
+   *
+   * @group agg_func
+   * @since 1.13.0
+   */
+  def countDistinct(colName: String, colNames: String*): Column =
+    count_distinct(col(colName), colNames.map(Column.apply): _*)
+
+  /**
+   * Returns either the number of non-NULL distinct records for the specified columns,
+   * or the total number of the distinct records. An alias of count_distinct.
+   *
+   * @group agg_func
+   * @since 1.13.0
+   */
+  def countDistinct(expr: Column, exprs: Column*): Column =
+    count_distinct(expr, exprs: _*)
+
   /**
    * Returns either the number of non-NULL distinct records for the specified columns,
    * or the total number of the distinct records.

src/test/java/com/snowflake/snowpark_test/JavaFunctionSuite.java

@@ -67,6 +67,14 @@ public void count() {
 
     Row[] expected2 = {Row.create(5)};
     checkAnswer(df.select(Functions.count_distinct(df.col("A"), df.col("b"))), expected2);
+
+    checkAnswer(df.select(Functions.countDistinct(df.col("A"))), expected1);
+
+    checkAnswer(df.select(Functions.countDistinct(df.col("A"), df.col("b"))), expected2);
+
+    checkAnswer(df.select(Functions.countDistinct("A")), expected1);
+
+    checkAnswer(df.select(Functions.countDistinct("A", "b")), expected2);
   }
 
   @Test

src/test/scala/com/snowflake/snowpark_test/FunctionSuite.scala

@@ -40,6 +40,10 @@ trait FunctionSuite extends TestData {
     checkAnswer(duplicatedNumbers.select(count(col("A"))), Seq(Row(5)))
 
     checkAnswer(duplicatedNumbers.select(count_distinct(col("A"))), Seq(Row(3)))
+
+    checkAnswer(duplicatedNumbers.select(countDistinct(col("A"))), Seq(Row(3)))
+
+    checkAnswer(duplicatedNumbers.select(countDistinct("A")), Seq(Row(3)))
   }
 
   test("covariance") {