Polars SQL 数组函数完全指南ARRAY_AGG、ARRAY_CONTAINS、ARRAY_GET 与 UNNEST 等 13 个函数的用法与源码解析【免费下载链接】polarsExtremely fast Query Engine for DataFrames, written in Rust项目地址: https://gitcode.com/GitHub_Trending/po/polarsPolars 内置 SQL 引擎在df.sql()中提供了 13 个面向数组Array/List的函数覆盖数组的聚合构造、元素访问、统计计算、字符串拼接、去重与行展开unnest等全部常见场景。本文以官方 API 参考文档 py-polars/docs/source/reference/sql/functions/array.rst 为主体骨架结合 crates/polars-sql/src/functions.rs 的底层实现与 crates/polars-sql/tests/simple_exprs.rs 的测试用例逐一讲解每个函数的语义、语法、返回类型与易错点读完即可在 DataFrame 上直接写出可运行的数组型 SQL 查询。一、准备工作在 Polars 中执行 SQL这些数组函数统一通过 DataFrame 的sql()方法调用SQL 的FROM指向当前 DataFramePolars 约定用FROM self引用调用方import polars as pl df pl.DataFrame({foo: [[1, 2], [4, 3, 2]]}) df.sql(SELECT ARRAY_LENGTH(foo) AS n FROM self)函数名大小写不敏感array_agg与ARRAY_AGG等价。SQL 引擎对数组类型做了两种区分List变长数组list[i64]、list[str]等普通 DataFrame 列最常见的数组形态Array定长数组如pl.Array(pl.Float64, 2)用于ARRAY_INNER_PRODUCT等要求等宽定长输入的场景。从源码看所有数组函数在 SQL 侧统一注册于 crates/polars-sql/src/functions.rs 的函数名表中并最终映射为 Polars List/Array 表达式 API如list().len()、list().max()因此 SQL 层函数与 Polars 原生表达式能力完全对齐。二、函数总览函数说明返回类型典型ARRAY_AGG将列/表达式聚合成数组等价于 implodelist[...]ARRAY_CONTAINS判断数组是否包含指定值boolARRAY_GET取数组中指定下标的值SQL 从 1 开始元素类型ARRAY_INNER_PRODUCT两个定长数组的内积别名ARRAY_DOT_PRODUCT数值ARRAY_LENGTH返回数组长度u32ARRAY_LOWER返回数组下界最小值元素类型ARRAY_MEAN返回数组均值f64ARRAY_REVERSE反转数组元素顺序list[...]ARRAY_SUM返回数组元素之和元素类型ARRAY_TO_STRING将所有元素拼接为字符串strARRAY_UNIQUE返回去重后的数组list[...]ARRAY_UPPER返回数组上界最大值元素类型UNNEST将数组列展开为多行元素类型三、ARRAY_AGG聚合构造数组ARRAY_AGG(expr)将一列的多个值聚合为一个数组语义上等价于 Polars 表达式的implode。这是分组聚合、嵌套结构构建中最常用的数组函数。3.1 基本用法与内联子句ARRAY_AGG支持可选的ORDER BY与LIMIT子句直接写在参数括号内df pl.DataFrame({foo: [1, 2, 3], bar: [4, 5, 6]}) df.sql( SELECT ARRAY_AGG(foo ORDER BY foo DESC) AS arr_foo, ARRAY_AGG(bar LIMIT 2) AS arr_bar FROM self ) # shape: (1, 2) # ┌───────────┬───────────┐ # │ arr_foo ┆ arr_bar │ # │ --- ┆ --- │ # │ list[i64] ┆ list[i64] │ # ╞═══════════╪═══════════╡ # │ [3, 2, 1] ┆ [4, 5] │ # └───────────┴───────────┘ARRAY_AGG(foo ORDER BY foo DESC)先按foo降序排序再聚合得到[3, 2, 1]ARRAY_AGG(bar LIMIT 2)只取前两行聚合得到[4, 5]。3.2 源码实现内联子句如何组合底层实现位于 crates/polars-sql/src/functions.rs 的visit_arr_agg以及 apply_aggregate_clausesORDER BY先对表达式执行sort_by排序再聚合LIMIT对排序后的结果先执行head(n)截断再聚合DISTINCT先unique_stable()去重若同时有ORDER BY则对去重结果排序最终统一调用base.implode(true)收尾成数组。测试用例 crates/polars-sql/tests/simple_exprs.rs 验证了这些组合其中SELECT ARRAY_AGG(a ORDER BY b LIMIT 2) FROM df被断言等价于col(a).sort_by(...).head(2).implode(true)——即排序、截断、聚合按此顺序执行。此外ARRAY_AGG只接受恰好一个参数多传会报 SQL 语法错误。四、ARRAY_CONTAINS判断元素是否存在ARRAY_CONTAINS(array, value)逐行判断数组是否包含给定值返回布尔列df pl.DataFrame({foo: [[1, 2], [4, 3]]}) df.sql(SELECT foo, ARRAY_CONTAINS(foo, 2) AS has_two FROM self) # shape: (2, 2) # ┌───────────┬─────────┐ # │ foo ┆ has_two │ # │ --- ┆ --- │ # │ list[i64] ┆ bool │ # ╞═══════════╪═════════╡ # │ [1, 2] ┆ true │ # │ [4, 3] ┆ false │ # └───────────┴─────────┘底层实现为e.list().contains(s, true)functions.rs即 Polars 的list.contains表达式第二个参数true表示严格匹配。五、ARRAY_GET按下标取元素注意是 1 基索引ARRAY_GET(array, index)返回数组中指定位置的元素。与编程语言常见的 0 基索引不同SQL 语义下索引从 1 开始越界时返回null而非报错df pl.DataFrame( { foo: [[1, 2], [4, 3, 2]], bar: [[6, 7], [8, 9, 10]] } ) df.sql( SELECT foo, bar, ARRAY_GET(foo, 1) AS foo_at_1, ARRAY_GET(bar, 3) AS bar_at_2 FROM self ) # shape: (2, 4) # ┌───────────┬────────────┬──────────┬──────────┐ # │ foo ┆ bar ┆ foo_at_1 ┆ bar_at_2 │ # │ --- ┆ --- ┆ --- ┆ --- │ # │ list[i64] ┆ list[i64] ┆ i64 ┆ i64 │ # ╞═══════════╪════════════╪══════════╪══════════╡ # │ [1, 2] ┆ [6, 7] ┆ 1 ┆ null │ # │ [4, 3, 2] ┆ [8, 9, 10] ┆ 4 ┆ 10 │ # └───────────┴────────────┴──────────┴──────────┘ARRAY_GET(foo, 1)取第一个元素首行得到1ARRAY_GET(bar, 3)取第三个元素首行[6, 7]只有两个元素越界返回null第二行[8, 9, 10]返回10。源码中专门有一段注释说明这一点functions.rs 通过adjust_one_indexed_param(idx, true)将 SQL 的 1 基索引转换为 Polars 表达式内部使用的 0 基索引再调用list().get(idx, true)第二个参数同样表示越界/空值时宽容返回 null。六、ARRAY_INNER_PRODUCT定长数组内积ARRAY_INNER_PRODUCT(lhs, rhs)计算两个等宽定长数组的内积点积ARRAY_DOT_PRODUCT是其别名。要点两侧元素数据类型会先 cast 到公共超类型该超类型必须是整数、Float32或Float64输入可以是定长 Array 表达式也可以是已知宽度的 SQL 数组字面量字面量会被解释为标量定长数组并与另一侧输入做广播其他具有变长 List 类型的表达式不会被隐式转换某一坐标处任一元素为 null 时该坐标不参与求和非 null 行若无任何有效坐标对则结果为0若某行任一输入 Array 本身为 null则该行结果为null。6.1 两个数组列做内积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 │ # └───────┘第 1 行1.0*10.0 2.0*20.0 50.0第 2 行3.0*30.0 4.0*40.0 250.0。6.2 数组字面量作为标量广播df.sql(SELECT ARRAY_INNER_PRODUCT(lhs, [10.0, 20.0]) AS dot FROM self) # shape: (2, 1) # ┌───────┐ # │ dot │ # │ --- │ # │ f64 │ # ╞═══════╡ # │ 50.0 │ # │ 110.0 │ # └───────┘这里[10.0, 20.0]被当作标量定长数组对每一行与lhs计算内积第 1 行1.0*10.0 2.0*20.0 50.0第 2 行3.0*10.0 4.0*20.0 110.0。6.3 源码实现实现位于 visit_array_inner_product 与 parse_array_inner_product_arg只有直接书写的 SQL 数组字面量SQLExpr::Array会被解析为定长标量数组用于arr().dot()广播其他表达式走普通parse_sql_arg路径保持 List 形态不做隐式转换。函数注册表同时收录了array_inner_product与array_dot_product两个名称functions.rs。七、ARRAY_LENGTH / ARRAY_LOWER / ARRAY_UPPER长度与极值三个一元函数分别返回数组长度、最小值、最大值是聚合类统计的基础df pl.DataFrame({foo: [[1, 2], [4, 3, 2]]}) df.sql(SELECT foo, ARRAY_LENGTH(foo) AS n_elems FROM self) # shape: (2, 2) # ┌───────────┬─────────┐ # │ foo ┆ n_elems │ # │ --- ┆ --- │ # │ list[i64] ┆ u32 │ # ╞═══════════╪═════════╡ # │ [1, 2] ┆ 2 │ # │ [4, 3, 2] ┆ 3 │ # └───────────┴─────────┘df pl.DataFrame({foo: [[1, 2], [4, -2, 8]]}) df.sql(SELECT foo, ARRAY_LOWER(foo) AS min_elem FROM self) # shape: (2, 2) # ┌────────────┬──────────┐ # │ foo ┆ min_elem │ # │ --- ┆ --- │ # │ list[i64] ┆ i64 │ # ╞════════════╪══════════╡ # │ [1, 2] ┆ 1 │ # │ [4, -2, 8] ┆ -2 │ # └────────────┴──────────┘df pl.DataFrame({foo: [[5, 0], [4, 8, -2]]}) df.sql(SELECT foo, ARRAY_UPPER(foo) AS max_elem FROM self) # shape: (2, 2) # ┌────────────┬──────────┐ # │ foo ┆ max_elem │ # │ --- ┆ --- │ # │ list[i64] ┆ i64 │ # ╞════════════╪══════════╡ # │ [5, 0] ┆ 5 │ # │ [4, 8, -2] ┆ 8 │ # └────────────┴──────────┘源码中这三个函数分别映射为list().len()、list().min()、list().max()functions.rs。值得注意的是文档注释明确指出ARRAY_LOWER等价于array_min、ARRAY_UPPER等价于array_maxfunctions.rs因此它们返回的是数组内元素的数值极值而非某些 SQL 方言中“数组下界下标”的含义。八、ARRAY_MEAN / ARRAY_SUM聚合统计对数组元素做均值与求和df pl.DataFrame({foo: [[1, 2], [4, 3, -1]]}) df.sql(SELECT foo, ARRAY_MEAN(foo) AS foo_mean FROM self) # shape: (2, 2) # ┌────────────┬──────────┐ # │ foo ┆ foo_mean │ # │ --- ┆ --- │ # │ list[i64] ┆ f64 │ # ╞════════════╪══════════╡ # │ [1, 2] ┆ 1.5 │ # │ [4, 3, -1] ┆ 2.0 │ # └────────────┴──────────┘df pl.DataFrame({foo: [[1, -2], [10, 3, -2]]}) df.sql(SELECT foo, ARRAY_SUM(foo) AS foo_sum FROM self) # shape: (2, 2) # ┌─────────────┬─────────┐ # │ foo ┆ foo_sum │ # │ --- ┆ --- │ # │ list[i64] ┆ i64 │ # ╞═════════════╪═════════╡ # │ [1, -2] ┆ -1 │ # │ [10, 3, -2] ┆ 11 │ # └─────────────┴─────────┘注意返回类型差异ARRAY_MEAN恒为浮点f64ARRAY_SUM保持元素类型整型数组求和为i64。源码对应list().mean()与list().sum()functions.rs。九、ARRAY_REVERSE反转数组ARRAY_REVERSE(array)逐行反转元素顺序多列可同时处理df pl.DataFrame( { foo: [[1, 2], [4, 3, 2]], bar: [[6, 7], [8, 9, 10]] } ) df.sql( SELECT foo, ARRAY_REVERSE(foo) AS oof, ARRAY_REVERSE(bar) AS rab FROM self ) # shape: (2, 3) # ┌───────────┬───────────┬────────────┐ # │ foo ┆ oof ┆ rab │ # │ --- ┆ --- ┆ --- │ # │ list[i64] ┆ list[i64] ┆ list[i64] │ # ╞═══════════╪═══════════╪════════════╡ # │ [1, 2] ┆ [2, 1] ┆ [7, 6] │ # │ [4, 3, 2] ┆ [2, 3, 4] ┆ [10, 9, 8] │ # └───────────┴───────────┴────────────┘实现上是逐元素表达式求值list().eval(element().reverse())functions.rs。十、ARRAY_TO_STRING数组转字符串ARRAY_TO_STRING(array, separator)将数组所有元素拼接为一个字符串分隔符由第二个参数指定df pl.DataFrame( { foo: [[a, b], [c, d, e]], bar: [[8, None, 8], [3, 2, 1, 0]], } ) df.sql( SELECT ARRAY_TO_STRING(foo,:) AS s_foo, ARRAY_TO_STRING(bar,:) AS s_bar FROM self ) # shape: (2, 2) # ┌───────┬─────────┐ # │ s_foo ┆ s_bar │ # │ --- ┆ --- │ # │ str ┆ str │ # ╞═══════╪═════════╡ # │ a:b ┆ 8:8 │ # │ c:d:e ┆ 3:2:1:0 │ # └───────┴─────────┘两个值得注意的行为在官方文档示例中直接可见元素会被 cast 为字符串整型数组[8, None, 8]也能拼接null 元素默认被忽略[8, None, 8]拼接结果是8:8中间的 null 被跳过而不是输出空段。底层实现 visit_arr_to_string 支持2 或 3 个参数两参数形式等价于list().join(separator, true)true即忽略 null三参数形式在启用list_evalfeature 时第三个参数作为 null 值的替换文本——若替换文本非空则先通过list().eval(element().fill_null(替换值))填充 null 再拼接。错误处理也很明确参数个数不是 2 或 3 时抛出ARRAY_TO_STRING expects 2-3 arguments的语法错误第三参数非字符串字面量时抛出invalid null value for ARRAY_TO_STRING。十一、ARRAY_UNIQUE数组内去重ARRAY_UNIQUE(array)返回去重后的数组保留首次出现的顺序df pl.DataFrame({foo: [[a, b], [b, b, e]]}) df.sql(SELECT ARRAY_UNIQUE(foo) AS foo_unique FROM self) # shape: (2, 1) # ┌────────────┐ # │ foo_unique │ # │ --- │ # │ list[str] │ # ╞════════════╡ # │ [a, b] │ # │ [e, b] │ # └────────────┘第二行[b, b, e]去重后得到[e, b]——注意顺序是“先出现者优先”e排在b前面说明实现使用的是稳定去重list().eval(element().unique_stable())functions.rs。这与ARRAY_AGG(DISTINCT ...)的去重语义一致。十二、UNNEST数组列展开为多行UNNEST(column)将数组列拆开、每个元素各占一行是“长表化”explode的核心操作。多个数组列同时展开时按位置对齐df pl.DataFrame( { foo: [[a, b, c], [d, e]], bar: [[6, 7, 8], [9, 10]] } ) df.sql( SELECT UNNEST(foo) AS f, UNNEST(bar) AS b FROM self ) # shape: (5, 2) # ┌─────┬─────┐ # │ f ┆ b │ # │ --- ┆ --- │ # │ str ┆ i64 │ # ╞═════╪═════╡ # │ a ┆ 6 │ # │ b ┆ 7 │ # │ c ┆ 8 │ # │ d ┆ 9 │ # │ e ┆ 10 │ # └─────┴─────┘foo两行3 个、2 个元素与bar两行3 个、2 个元素同时展开得到 5 行。在函数形式中UNNEST 映射为 Explode 分支e.explode(ExplodeOptions { empty_as_null: true, keep_nulls: true })即空数组按 null 处理、null 元素保留。此外 UNNEST 还有第二种形态——表函数形式CROSS JOIN UNNEST(col) AS t(alias)此时由 crates/polars-sql/src/context.rs 的process_unnest_lateral按横向lateral操作处理要求必须带表别名且暂不支持WITH ORDINALITY/WITH OFFSET会报UNNEST tables do not (yet) support WITH ORDINALITY|OFFSET。两种形态可在同一查询中组合例如测试用例中的SELECT UNNEST(ARRAY_AGG(DISTINCT a)) FROM dfsimple_exprs.rs先用ARRAY_AGG(DISTINCT ...)去重聚合再展开等价于取去重后的全部值。十三、组合实战聚合、展开与条件判断数组函数可以自由嵌套组合。以下基于 crates/polars-sql/tests/simple_exprs.rs 的测试模式演示 UNNEST 结果参与CASE WHEN判断的写法df pl.DataFrame( { id: [1, 2], list_a: [[x, y], [z]], list_b: [[6000, 3000], [1000]], } ) df.sql( SELECT id, unnest(list_a) AS list_a, unnest(list_b) AS list_b, CASE WHEN unnest(list_b) 5000 THEN High WHEN unnest(list_b) 2500 THEN Medium ELSE Low END AS tier FROM self )该查询将list_b展开后按数值分档是“嵌套结构 → 扁平明细 → 业务分桶”的典型链路也可与WHERE、GROUP BY、窗口函数等 SQL 子句配合使用。十四、易错点与最佳实践ARRAY_GET 是 1 基索引ARRAY_GET(foo, 1)取第一个元素与 Python 列表的foo[0]不同越界返回null而非报错functions.rs。ARRAY_INNER_PRODUCT 只接受定长 Array 或等宽字面量变长 List 不会被隐式转换字面量作为标量广播functions.rs。ARRAY_TO_STRING 忽略 null需要占位文本时使用三参数形式需list_evalfeature。ARRAY_UNIQUE 保持首次出现顺序若需排序后的去重结果可先ARRAY_AGG(... ORDER BY ...)再处理或对结果另行排序。ARRAY_LOWER/ARRAY_UPPER 是极值而非下标语义与min/max等价functions.rs。UNNEST 表函数必须带别名且不支持WITH ORDINALITYcontext.rs。语法错误信息明确如ARRAY_AGG must have exactly one argument、ARRAY_TO_STRING expects 2-3 arguments、LIMIT in ARRAY_AGG must be a positive integer遇到后按提示修正即可。十五、小结本文以 py-polars/docs/source/reference/sql/functions/array.rst 为准绳完整覆盖了 Polars SQL 引擎的全部 13 个数组函数聚合构造ARRAY_AGG、元素访问ARRAY_CONTAINS、ARRAY_GET、统计ARRAY_LENGTH/ARRAY_LOWER/ARRAY_UPPER/ARRAY_MEAN/ARRAY_SUM、变换ARRAY_REVERSE、ARRAY_UNIQUE、ARRAY_TO_STRING、定长内积ARRAY_INNER_PRODUCT/ARRAY_DOT_PRODUCT与行展开UNNEST。每个函数都给出了可直接运行的示例、返回类型与底层表达式映射便于在 crates/polars-sql/src/functions.rs 与 crates/polars-sql/tests 中继续核对行为细节。掌握这组函数后你可以在df.sql()中直接完成嵌套数组的构建、解析、统计与扁平化将 SQL 的声明式表达力与 Polars 的向量化执行结合起来。【免费下载链接】polarsExtremely fast Query Engine for DataFrames, written in Rust项目地址: https://gitcode.com/GitHub_Trending/po/polars创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考