[SPARK-58965][SQL][PYTHON] Add gcd and lcm math functions

python/docs/source/reference/pyspark.sql/functions.rst

@@ -113,9 +113,11 @@ Mathematical Functions
     expm1
     factorial
     floor
+    gcd
     greatest
     hex
     hypot
+    lcm
     least
     ln
     log

python/pyspark/sql/connect/functions/builtin.py

@@ -709,6 +709,20 @@ def factorial(col: "ColumnOrName") -> Column:
 factorial.__doc__ = pysparkfuncs.factorial.__doc__
 
 
+def gcd(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+    return _invoke_function_over_columns("gcd", col1, col2)
+
+
+gcd.__doc__ = pysparkfuncs.gcd.__doc__
+
+
+def lcm(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+    return _invoke_function_over_columns("lcm", col1, col2)
+
+
+lcm.__doc__ = pysparkfuncs.lcm.__doc__
+
+
 def floor(col: "ColumnOrName", scale: Optional[Union[Column, int]] = None) -> Column:
     if scale is None:
         return _invoke_function_over_columns("floor", col)

python/pyspark/sql/functions/__init__.py

@@ -80,9 +80,11 @@
     "expm1",
     "factorial",
     "floor",
+    "gcd",
     "greatest",
     "hex",
     "hypot",
+    "lcm",
     "least",
     "ln",
     "log",

python/pyspark/sql/functions/builtin.py

@@ -9297,6 +9297,92 @@ def factorial(col: "ColumnOrName") -> Column:
     return _invoke_function_over_columns("factorial", col)
 
 
+@_try_remote_functions
+def gcd(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+    """
+    Computes the greatest common divisor of the two given values.
+    The result is non-negative, and is 0 when both values are 0.
+
+    .. versionadded:: 4.4.0
+
+    Parameters
+    ----------
+    col1 : :class:`~pyspark.sql.Column` or str
+        the first value.
+        A column that evaluates to an integral.
+    col2 : :class:`~pyspark.sql.Column` or str
+        the second value.
+        A column that evaluates to an integral.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        greatest common divisor of the given values.
+        Returns a column that evaluates to a long.
+
+    See Also
+    --------
+    :meth:`pyspark.sql.functions.lcm`
+
+    Examples
+    --------
+    >>> from pyspark.sql import functions as sf
+    >>> spark.range(1, 5).select("*", sf.gcd(sf.lit(12), 'id')).show()
+    +---+-----------+
+    | id|gcd(12, id)|
+    +---+-----------+
+    |  1|          1|
+    |  2|          2|
+    |  3|          3|
+    |  4|          4|
+    +---+-----------+
+    """
+    return _invoke_function_over_columns("gcd", col1, col2)
+
+
+@_try_remote_functions
+def lcm(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
+    """
+    Computes the least common multiple of the two given values.
+    The result is non-negative, and is 0 when either value is 0.
+
+    .. versionadded:: 4.4.0
+
+    Parameters
+    ----------
+    col1 : :class:`~pyspark.sql.Column` or str
+        the first value.
+        A column that evaluates to an integral.
+    col2 : :class:`~pyspark.sql.Column` or str
+        the second value.
+        A column that evaluates to an integral.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        least common multiple of the given values.
+        Returns a column that evaluates to a long.
+
+    See Also
+    --------
+    :meth:`pyspark.sql.functions.gcd`
+
+    Examples
+    --------
+    >>> from pyspark.sql import functions as sf
+    >>> spark.range(1, 5).select("*", sf.lcm(sf.lit(6), 'id')).show()
+    +---+----------+
+    | id|lcm(6, id)|
+    +---+----------+
+    |  1|         6|
+    |  2|         6|
+    |  3|         6|
+    |  4|        12|
+    +---+----------+
+    """
+    return _invoke_function_over_columns("lcm", col1, col2)
+
+
 @_try_remote_functions
 def lag(col: "ColumnOrName", offset: int = 1, default: Optional[Any] = None) -> Column:
     """

sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/MathUtils.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.util
 
+import java.{lang => jl}
+
 import org.apache.spark.QueryContext
 import org.apache.spark.sql.errors.ExecutionErrors
 
@@ -123,6 +125,54 @@ object MathUtils {
     if (r < 0) (r + n) % n else r
   }
 
+  // Greatest common divisor of two longs, computed with the Euclidean algorithm. The result is
+  // always non-negative, and `gcd(0, 0)` is 0. The only unrepresentable result is `-Long.MinValue`,
+  // reached by `(0, x)`, `(x, 0)` and `(x, x)` for `x == Long.MinValue`; as elsewhere in Spark that
+  // overflow raises under ANSI mode and yields null otherwise.
+  def gcd(a: Long, b: Long, ansiEnabled: Boolean, context: QueryContext): jl.Long = {
+    var x = a
+    var y = b
+    while (y != 0) {
+      val remainder = x % y
+      x = y
+      y = remainder
+    }
+    // `x` carries the sign of the inputs, so take the absolute value to normalize the result.
+    if (x == Long.MinValue) {
+      overflowOrNull(ansiEnabled, context)
+    } else {
+      Math.abs(x)
+    }
+  }
+
+  // Least common multiple of two longs. Dividing by the greatest common divisor before multiplying
+  // keeps the intermediate product as small as possible, so only genuinely unrepresentable results
+  // overflow. The result is always non-negative, and is 0 when either input is 0.
+  def lcm(a: Long, b: Long, ansiEnabled: Boolean, context: QueryContext): jl.Long = {
+    if (a == 0 || b == 0) {
+      jl.Long.valueOf(0L)
+    } else {
+      val divisor = gcd(a, b, ansiEnabled, context)
+      if (divisor == null) {
+        null
+      } else {
+        try {
+          Math.multiplyExact(Math.absExact(a / divisor.longValue()), Math.absExact(b))
+        } catch {
+          case _: ArithmeticException => overflowOrNull(ansiEnabled, context)
+        }
+      }
+    }
+  }
+
+  private def overflowOrNull(ansiEnabled: Boolean, context: QueryContext): jl.Long = {
+    if (ansiEnabled) {
+      throw ExecutionErrors.arithmeticOverflowError("long overflow", context = context)
+    } else {
+      null
+    }
+  }
+
   // Casts a rounded double (the result of Math.ceil/Math.floor) to long, throwing an arithmetic
   // overflow error when the value cannot be represented as a long. NaN is passed through to the
   // JVM cast (which yields 0), matching the previous behavior. Shared by the eval and codegen

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

@@ -5573,6 +5573,21 @@ object functions {
    */
   def floor(columnName: String): Column = floor(Column(columnName))
 
+  /**
+   * Computes the greatest common divisor of the two given values. The result is non-negative, and
+   * is 0 when both values are 0.
+   *
+   * @param l
+   *   the first value. A column that evaluates to an integral.
+   * @param r
+   *   the second value. A column that evaluates to an integral.
+   * @group math_funcs
+   * @since 4.4.0
+   * @return
+   *   Returns a column that evaluates to a long.
+   */
+  def gcd(l: Column, r: Column): Column = Column.fn("gcd", l, r)
+
   /**
    * Returns the greatest value of the list of values, skipping null values. This function takes
    * at least 2 parameters. It will return null iff all parameters are null.
@@ -5743,6 +5758,21 @@ object functions {
    */
   def hypot(l: Double, rightName: String): Column = hypot(l, Column(rightName))
 
+  /**
+   * Computes the least common multiple of the two given values. The result is non-negative, and
+   * is 0 when either value is 0.
+   *
+   * @param l
+   *   the first value. A column that evaluates to an integral.
+   * @param r
+   *   the second value. A column that evaluates to an integral.
+   * @group math_funcs
+   * @since 4.4.0
+   * @return
+   *   Returns a column that evaluates to a long.
+   */
+  def lcm(l: Column, r: Column): Column = Column.fn("lcm", l, r)
+
   /**
    * Returns the least value of the list of values, skipping null values. This function takes at
    * least 2 parameters. It will return null iff all parameters are null.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -476,8 +476,10 @@ object FunctionRegistry {
     expression[Expm1]("expm1"),
     expressionBuilder("floor", FloorExpressionBuilder),
     expression[Factorial]("factorial"),
+    expression[Gcd]("gcd"),
     expression[Hex]("hex"),
     expression[Hypot]("hypot"),
+    expression[Lcm]("lcm"),
     expression[Logarithm]("log"),
     expression[Log10]("log10"),
     expression[Log1p]("log1p"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala

@@ -762,6 +762,138 @@ case class Factorial(child: Expression)
     copy(child = newChild)
 }
 
+@ExpressionDescription(
+  usage = "_FUNC_(expr1, expr2) - Returns the greatest common divisor of `expr1` and `expr2`.",
+  arguments = """
+    Arguments:
+      * expr1 - The first value. An expression that evaluates to an integral number.
+      * expr2 - The second value. An expression that evaluates to an integral number.
+  """,
+  examples = """
+    Examples:
+      > SELECT _FUNC_(24, 36);
+       12
+      > SELECT _FUNC_(-24, 36);
+       12
+      > SELECT _FUNC_(0, 0);
+       0
+  """,
+  since = "4.4.0",
+  group = "math_funcs")
+case class Gcd(
+    left: Expression,
+    right: Expression,
+    ansiEnabled: Boolean = SQLConf.get.ansiEnabled)
+  extends BinaryExpression with ImplicitCastInputTypes with SupportQueryContext {
+  override def nullIntolerant: Boolean = true
+
+  def this(left: Expression, right: Expression) =
+    this(left, right, ansiEnabled = SQLConf.get.ansiEnabled)
+
+  override def inputTypes: Seq[DataType] = Seq(LongType, LongType)
+
+  override def dataType: DataType = LongType
+
+  // Overflow yields null in non-ANSI mode.
+  override def nullable: Boolean = true
+
+  override def initQueryContext(): Option[QueryContext] = if (ansiEnabled) {
+    Some(origin.context)
+  } else {
+    None
+  }
+
+  protected override def nullSafeEval(left: Any, right: Any): Any = {
+    MathUtils.gcd(
+      left.asInstanceOf[Long], right.asInstanceOf[Long], ansiEnabled, getContextOrNull())
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val errorContext = getContextOrNullCode(ctx, ansiEnabled)
+    val result = ctx.freshName("gcd")
+    nullSafeCodeGen(ctx, ev, (leftValue, rightValue) => {
+      s"""
+        java.lang.Long $result = org.apache.spark.sql.catalyst.util.MathUtils.gcd(
+          $leftValue, $rightValue, $ansiEnabled, $errorContext);
+        if ($result == null) {
+          ${ev.isNull} = true;
+        } else {
+          ${ev.value} = $result;
+        }
+      """
+    })
+  }
+
+  override protected def withNewChildrenInternal(
+      newLeft: Expression, newRight: Expression): Gcd = copy(left = newLeft, right = newRight)
+}
+
+@ExpressionDescription(
+  usage = "_FUNC_(expr1, expr2) - Returns the least common multiple of `expr1` and `expr2`.",
+  arguments = """
+    Arguments:
+      * expr1 - The first value. An expression that evaluates to an integral number.
+      * expr2 - The second value. An expression that evaluates to an integral number.
+  """,
+  examples = """
+    Examples:
+      > SELECT _FUNC_(4, 6);
+       12
+      > SELECT _FUNC_(-4, 6);
+       12
+      > SELECT _FUNC_(0, 5);
+       0
+  """,
+  since = "4.4.0",
+  group = "math_funcs")
+case class Lcm(
+    left: Expression,
+    right: Expression,
+    ansiEnabled: Boolean = SQLConf.get.ansiEnabled)
+  extends BinaryExpression with ImplicitCastInputTypes with SupportQueryContext {
+  override def nullIntolerant: Boolean = true
+
+  def this(left: Expression, right: Expression) =
+    this(left, right, ansiEnabled = SQLConf.get.ansiEnabled)
+
+  override def inputTypes: Seq[DataType] = Seq(LongType, LongType)
+
+  override def dataType: DataType = LongType
+
+  // Overflow yields null in non-ANSI mode.
+  override def nullable: Boolean = true
+
+  override def initQueryContext(): Option[QueryContext] = if (ansiEnabled) {
+    Some(origin.context)
+  } else {
+    None
+  }
+
+  protected override def nullSafeEval(left: Any, right: Any): Any = {
+    MathUtils.lcm(
+      left.asInstanceOf[Long], right.asInstanceOf[Long], ansiEnabled, getContextOrNull())
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val errorContext = getContextOrNullCode(ctx, ansiEnabled)
+    val result = ctx.freshName("lcm")
+    nullSafeCodeGen(ctx, ev, (leftValue, rightValue) => {
+      s"""
+        java.lang.Long $result = org.apache.spark.sql.catalyst.util.MathUtils.lcm(
+          $leftValue, $rightValue, $ansiEnabled, $errorContext);
+        if ($result == null) {
+          ${ev.isNull} = true;
+        } else {
+          ${ev.value} = $result;
+        }
+      """
+    })
+  }
+
+  override protected def withNewChildrenInternal(
+      newLeft: Expression, newRight: Expression): Lcm = copy(left = newLeft, right = newRight)
+}
+
 @ExpressionDescription(
   usage = "_FUNC_(expr) - Returns the natural logarithm (base e) of `expr`.",
   arguments = """