From dfc1c782518e301ec8c4db4e76c08bffaeb39c21 Mon Sep 17 00:00:00 2001
From: Luis Bravo Zuniga <luis.bravozuniga@snowflake.com>
Date: Wed, 24 Apr 2024 14:22:05 -0600
Subject: [PATCH] Add com.snowflake.snowpark.functions.countDistinct (#102)

Add com.snowflake.snowpark.functions.countDistinct
---
 .../snowflake/snowpark_java/Functions.java    | 31 +++++++++++++++++++
 .../com/snowflake/snowpark/functions.scala    | 20 ++++++++++++
 .../snowpark_test/JavaFunctionSuite.java      |  8 +++++
 .../snowpark_test/FunctionSuite.scala         |  4 +++
 4 files changed, 63 insertions(+)

diff --git a/src/main/java/com/snowflake/snowpark_java/Functions.java b/src/main/java/com/snowflake/snowpark_java/Functions.java
index 416d8c1b..0515164a 100644
--- a/src/main/java/com/snowflake/snowpark_java/Functions.java
+++ b/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.
diff --git a/src/main/scala/com/snowflake/snowpark/functions.scala b/src/main/scala/com/snowflake/snowpark/functions.scala
index ee2676da..0f48d01f 100644
--- a/src/main/scala/com/snowflake/snowpark/functions.scala
+++ b/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.
diff --git a/src/test/java/com/snowflake/snowpark_test/JavaFunctionSuite.java b/src/test/java/com/snowflake/snowpark_test/JavaFunctionSuite.java
index 42253362..6ee298d3 100644
--- a/src/test/java/com/snowflake/snowpark_test/JavaFunctionSuite.java
+++ b/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
diff --git a/src/test/scala/com/snowflake/snowpark_test/FunctionSuite.scala b/src/test/scala/com/snowflake/snowpark_test/FunctionSuite.scala
index 38cad6c5..e473de12 100644
--- a/src/test/scala/com/snowflake/snowpark_test/FunctionSuite.scala
+++ b/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") {