feat: Expose fixed-array inner product in SQL

crates/polars-sql/src/functions.rs

@@ -2,7 +2,7 @@ use std::ops::{Add, Sub};
 
 use polars_core::chunked_array::ops::{FillNullStrategy, SortMultipleOptions, SortOptions};
 use polars_core::prelude::{
-    DataType, ExplodeOptions, PolarsResult, QuantileMethod, Schema, TimeUnit, polars_bail,
+    DataType, ExplodeOptions, PolarsResult, QuantileMethod, Scalar, Schema, TimeUnit, polars_bail,
     polars_err,
 };
 use polars_lazy::dsl::Expr;
@@ -27,7 +27,9 @@ use sqlparser::ast::{
 use sqlparser::tokenizer::Span;
 
 use crate::SQLContext;
-use crate::sql_expr::{adjust_one_indexed_param, parse_extract_date_part, parse_sql_expr};
+use crate::sql_expr::{
+    adjust_one_indexed_param, parse_extract_date_part, parse_sql_array, parse_sql_expr,
+};
 
 pub(crate) struct SQLFunctionVisitor<'a> {
     pub(crate) func: &'a SQLFunction,
@@ -730,6 +732,12 @@ pub(crate) enum PolarsSQLFunctions {
     /// SELECT ARRAY_CONTAINS(col1, 'foo') FROM df;
     /// ```
     ArrayContains,
+    /// SQL 'array_inner_product' function (also known as `array_dot_product`).
+    /// Returns the inner product of two fixed-size arrays.
+    /// ```sql
+    /// SELECT ARRAY_INNER_PRODUCT(col1, col2) FROM df;
+    /// ```
+    ArrayInnerProduct,
     /// SQL 'unnest' function.
     /// Unnest/explodes an array column into multiple rows.
     /// ```sql
@@ -809,7 +817,9 @@ impl PolarsSQLFunctions {
             "acos",
             "acosd",
             "array_contains",
+            "array_dot_product",
             "array_get",
+            "array_inner_product",
             "array_length",
             "array_lower",
             "array_mean",
@@ -1057,6 +1067,7 @@ impl PolarsSQLFunctions {
             // ----
             "array_agg" => Self::ArrayAgg,
             "array_contains" => Self::ArrayContains,
+            "array_dot_product" | "array_inner_product" => Self::ArrayInnerProduct,
             "array_get" => Self::ArrayGet,
             "array_length" => Self::ArrayLength,
             "array_lower" => Self::ArrayMin,
@@ -1659,6 +1670,7 @@ impl SQLFunctionVisitor<'_> {
             // ----
             ArrayAgg => self.visit_arr_agg(),
             ArrayContains => self.visit_binary::<Expr>(|e, s| e.list().contains(s, true)),
+            ArrayInnerProduct => self.visit_array_inner_product(),
             ArrayGet => {
                 // note: SQL is 1-indexed, not 0-indexed
                 self.visit_binary(|e, idx: Expr| {
@@ -1988,10 +2000,38 @@ impl SQLFunctionVisitor<'_> {
     /// active `FILTER (WHERE …)` clause from the surrounding call.
     fn parse_sql_arg(&mut self, expr: &SQLExpr) -> PolarsResult<Expr> {
         let parsed = parse_sql_expr(expr, self.ctx, self.active_schema)?;
-        Ok(match &self.filter {
-            Some(pred) => parsed.filter(pred.clone()),
-            None => parsed,
-        })
+        Ok(self.apply_filter(parsed))
+    }
+
+    fn apply_filter(&self, expr: Expr) -> Expr {
+        match &self.filter {
+            Some(pred) => expr.filter(pred.clone()),
+            None => expr,
+        }
+    }
+
+    fn parse_array_inner_product_arg(&mut self, expr: &SQLExpr) -> PolarsResult<Expr> {
+        // Keep ordinary SQL arrays List-backed. Only direct literals in this
+        // function become scalar Arrays so native arr.dot can broadcast them.
+        let array_expr = match expr {
+            SQLExpr::Array(_) => expr,
+            SQLExpr::Nested(inner) => return self.parse_array_inner_product_arg(inner),
+            _ => return self.parse_sql_arg(expr),
+        };
+        let values = parse_sql_array(array_expr, self.ctx)?;
+        let width = values.len();
+        Ok(self.apply_filter(lit(Scalar::new_array(values, width))))
+    }
+
+    fn visit_array_inner_product(&mut self) -> PolarsResult<Expr> {
+        let args = extract_args(self.func)?;
+        match args.as_slice() {
+            [FunctionArgExpr::Expr(lhs), FunctionArgExpr::Expr(rhs)] => Ok(self
+                .parse_array_inner_product_arg(lhs)?
+                .arr()
+                .dot(self.parse_array_inner_product_arg(rhs)?)),
+            _ => self.not_supported_error(),
+        }
     }
 
     fn visit_unary(&mut self, f: impl Fn(Expr) -> Expr) -> PolarsResult<Expr> {

crates/polars-sql/src/keywords.rs

@@ -82,3 +82,16 @@ pub fn all_functions() -> Vec<&'static str> {
     functions.extend_from_slice(PolarsSQLFunctions::keywords());
     functions
 }
+
+#[cfg(test)]
+mod tests {
+    use super::all_functions;
+
+    #[test]
+    fn array_inner_product_aliases_are_discoverable() {
+        let functions = all_functions();
+        for name in ["array_inner_product", "array_dot_product"] {
+            assert!(functions.contains(&name));
+        }
+    }
+}

py-polars/docs/source/reference/sql/functions/array.rst

@@ -13,6 +13,8 @@ Array
      - Returns true if the array contains the value.
    * - :ref:`ARRAY_GET <array_get>`
      - Returns the value at the given index in the array.
+   * - :ref:`ARRAY_INNER_PRODUCT <array_inner_product>`
+     - Returns the inner product of two fixed-size arrays. Alias: ``ARRAY_DOT_PRODUCT``.
    * - :ref:`ARRAY_LENGTH <array_length>`
      - Returns the length of the array.
    * - :ref:`ARRAY_LOWER <array_lower>`
@@ -118,6 +120,61 @@ Returns the value at the given index in the array.
     # │ [4, 3, 2] ┆ [8, 9, 10] ┆ 4        ┆ 10       │
     # └───────────┴────────────┴──────────┴──────────┘
 
+.. _array_inner_product:
+
+ARRAY_INNER_PRODUCT
+-------------------
+Returns the inner product of two fixed-size arrays of equal width. Their inner
+data types are cast to a common supertype, which must be an integer, ``Float32``,
+or ``Float64``. ``ARRAY_DOT_PRODUCT`` is an alias.
+
+Inputs may be fixed-size Array expressions or direct, known-width SQL array
+literals. Direct literals are interpreted as scalar fixed-size Arrays and
+broadcast against the other input. Other expressions with variable-size List
+dtype are not implicitly converted.
+
+Coordinates where either array element is null do not contribute to the result.
+A non-null row with no valid coordinate pairs returns zero. If either input
+Array is null for a row, the result for that row is null.
+
+**Example:**
+
+.. code-block:: python
+
+    dtype = pl.Array(pl.Float64, 2)
+    df = pl.DataFrame(
+        {
+            "lhs": [[1.0, 2.0], [3.0, 4.0]],
+            "rhs": [[10.0, 20.0], [30.0, 40.0]],
+        },
+        schema={"lhs": dtype, "rhs": dtype},
+    )
+    df.sql("""
+      SELECT ARRAY_INNER_PRODUCT(lhs, rhs) AS dot FROM self
+    """)
+    # shape: (2, 1)
+    # ┌───────┐
+    # │ dot   │
+    # │ ---   │
+    # │ f64   │
+    # ╞═══════╡
+    # │ 50.0  │
+    # │ 250.0 │
+    # └───────┘
+
+    df.sql("""
+      SELECT ARRAY_INNER_PRODUCT(lhs, [10.0, 20.0]) AS dot FROM self
+    """)
+    # shape: (2, 1)
+    # ┌───────┐
+    # │ dot   │
+    # │ ---   │
+    # │ f64   │
+    # ╞═══════╡
+    # │ 50.0  │
+    # │ 110.0 │
+    # └───────┘
+
 .. _array_length:
 
 ARRAY_LENGTH

py-polars/tests/unit/sql/test_array.py

@@ -85,6 +85,138 @@ def test_array_literals() -> None:
         )
 
 
+@pytest.mark.parametrize("function_name", ["ARRAY_INNER_PRODUCT", "ARRAY_DOT_PRODUCT"])
+def test_array_inner_product(function_name: str) -> None:
+    df = pl.DataFrame(
+        {
+            "lhs": [[1, 2, 3], [1, None, 3], None, [None, None, None]],
+            "rhs": [
+                [4.0, 5.0, 6.0],
+                [4.0, 5.0, None],
+                [1.0, 2.0, 3.0],
+                [None, None, None],
+            ],
+        },
+        schema={
+            "lhs": pl.Array(pl.Int32, 3),
+            "rhs": pl.Array(pl.Float32, 3),
+        },
+    )
+
+    result = df.sql(f"SELECT {function_name}(lhs, rhs) AS dot FROM self")
+    expected = df.select(pl.col("lhs").arr.dot("rhs").alias("dot"))
+
+    assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    ("literal", "literal_on_left"),
+    [
+        ("[4.0, 5.0, 6.0]", False),
+        ("ARRAY[4.0, 5.0, 6.0]", True),
+        ("([4.0, 5.0, 6.0])", False),
+    ],
+)
+def test_array_inner_product_literal(literal: str, literal_on_left: bool) -> None:
+    df = pl.DataFrame(
+        {"arr": [[1, 2, 3], [4, 5, 6], None]},
+        schema={"arr": pl.Array(pl.Int32, 3)},
+    )
+    arguments = f"{literal}, arr" if literal_on_left else f"arr, {literal}"
+
+    result = df.sql(f"SELECT ARRAY_INNER_PRODUCT({arguments}) AS dot FROM self")
+    literal_expr = pl.lit([4.0, 5.0, 6.0], dtype=pl.Array(pl.Float64, 3))
+    expected = df.select(
+        (
+            literal_expr.arr.dot("arr")
+            if literal_on_left
+            else pl.col("arr").arr.dot(literal_expr)
+        ).alias("dot")
+    )
+
+    assert_frame_equal(result, expected)
+
+
+def test_array_inner_product_null_literal() -> None:
+    df = pl.DataFrame(
+        {"values": [[1, 2, 3], [4, None, 6], None]},
+        schema={"values": pl.Array(pl.Int32, 3)},
+    )
+
+    result = df.sql(
+        "SELECT ARRAY_INNER_PRODUCT(values, [NULL, NULL, NULL]) AS dot FROM self"
+    )
+
+    assert_frame_equal(
+        result,
+        pl.DataFrame({"dot": [0, 0, None]}, schema={"dot": pl.Int32}),
+    )
+
+
+def test_array_inner_product_literal_inherits_projection_height() -> None:
+    df = pl.DataFrame({"row": [1, 2, 3]})
+
+    result = df.sql("SELECT ARRAY_INNER_PRODUCT([1, 2], ARRAY[3, 4]) AS dot FROM self")
+
+    assert_frame_equal(result, pl.DataFrame({"dot": [11, 11, 11]}))
+
+
+def test_array_inner_product_non_array_receiver() -> None:
+    df = pl.DataFrame({"lhs": [[1, 2]]})
+
+    with pytest.raises(
+        pl.exceptions.InvalidOperationError,
+        match="expected Array datatype",
+    ):
+        df.sql("SELECT ARRAY_INNER_PRODUCT(lhs, [3, 4]) FROM self")
+
+
+def test_array_inner_product_uses_native_array_dot_plan() -> None:
+    df = pl.DataFrame(
+        {
+            "lhs": [[1, 2]],
+            "rhs": [[3, 4]],
+        },
+        schema={
+            "lhs": pl.Array(pl.Int64, 2),
+            "rhs": pl.Array(pl.Int64, 2),
+        },
+    )
+
+    with pl.SQLContext(df=df) as ctx:
+        plan = ctx.execute(
+            "SELECT ARRAY_INNER_PRODUCT(lhs, rhs) AS dot FROM df"
+        ).explain(optimized=False)
+
+    assert ".arr.dot([" in plan
+    assert ").arr.sum()" not in plan
+    assert " * " not in plan
+
+
+@pytest.mark.parametrize(
+    "arguments",
+    ["lhs", "lhs, rhs, lhs"],
+)
+def test_array_inner_product_arity(arguments: str) -> None:
+    df = pl.DataFrame(
+        {"lhs": [[1, 2]], "rhs": [[3, 4]]},
+        schema={"lhs": pl.Array(pl.Int64, 2), "rhs": pl.Array(pl.Int64, 2)},
+    )
+
+    with pytest.raises(SQLInterfaceError, match="no function matches"):
+        df.sql(f"SELECT ARRAY_INNER_PRODUCT({arguments}) FROM self")
+
+
+def test_array_inner_product_unequal_widths() -> None:
+    df = pl.DataFrame(
+        {"lhs": [[1, 2]], "rhs": [[3, 4, 5]]},
+        schema={"lhs": pl.Array(pl.Int64, 2), "rhs": pl.Array(pl.Int64, 3)},
+    )
+
+    with pytest.raises(pl.exceptions.ShapeError, match="equal array widths"):
+        df.sql("SELECT ARRAY_INNER_PRODUCT(lhs, rhs) FROM self")
+
+
 @pytest.mark.parametrize(
     ("array_index", "expected"),
     [