如何在slick-pg查询中使用PostgreSQL的array_length函数过滤行

PostgreSQL原生支持array_length函数，slick-pg也提供了对应的Slick操作封装，但我不清楚怎么把这个函数加到现有Slick查询里，以下是我的具体场景：

现有Slick查询代码 #

import DatabaseObjects.profile.api._
val randomFunction: Rep[Double] = SimpleFunction.nullary[Double]("random")
val query: Query[GeneratorV1, GeneratorV1Row, Seq] = GeneratorV1
  .sortBy(_ => randomFunction)
  .filter(_.rasterizationJobConfig +>> "activity" === activityTypeEnum)
  .filter(_.jobStatus === GeneratorStatuses.success)
  .take(N)

目标实现的SQL逻辑 #

我需要添加array_length过滤条件，最终生成的SQL类似：

SELECT * from generator_v1 WHERE rasterization_job_config->>'activity' = 'PLANTING'
  AND job_status = 'success' AND array_length(worker_job_ids, 1) = 2 LIMIT 5;

期望的Slick伪代码 #

大概逻辑如下（伪代码）：

import DatabaseObjects.profile.api._
val randomFunction: Rep[Double] = SimpleFunction.nullary[Double]("random")
val query: Query[GeneratorV1, GeneratorV1Row, Seq] = GeneratorV1
  .sortBy(_ => randomFunction)
  .filter(_.rasterizationJobConfig +>> "activity" === activityTypeEnum)
  .filter(_.jobStatus === GeneratorStatuses.success)
  .filter(/* 伪代码 */ _.jobIds.array_length > 0 && _.jobIds.array_length < N)
  .take(N)

我尝试过的失败代码 #

以下是我试过但没成功的实现：

import DatabaseObjects.profile.api._
val randomFunction: Rep[Double] = SimpleFunction.nullary[Double]("random")

def arrayLength(rep: Rep[List[Long]]): Rep[Boolean] = {
  val expression = SimpleExpression.unary[List[Long], Boolean] { (s, queryBuilder) =>
    queryBuilder.sqlBuilder += "array_length("
    queryBuilder.expr(s)
    queryBuilder.sqlBuilder += ", 1) = 2"
  }
  expression.apply(rep)
}

val query: Query[GeneratorV1, GeneratorV1Row, Seq] = GeneratorV1
  .sortBy(_ => randomFunction)
  .filter(_.rasterizationJobConfig +>> "activity" === activityTypeEnum)
  .filter(_.jobStatus === GeneratorStatuses.success)
  .filter(row => (row.workerJobIds, arrayLength(row.workerJobIds)))
  .take(N)

解决方案 #

方案1：直接使用slick-pg提供的arrayLength函数 #

如果你已经引入了slick-pg的数组扩展，直接用封装好的函数即可：

import DatabaseObjects.profile.api._
import com.github.tminglei.slickpg.ArrayExtensions._ // 导入slick-pg数组扩展

val randomFunction: Rep[Double] = SimpleFunction.nullary[Double]("random")
val query: Query[GeneratorV1, GeneratorV1Row, Seq] = GeneratorV1
  .sortBy(_ => randomFunction)
  .filter(_.rasterizationJobConfig +>> "activity" === activityTypeEnum)
  .filter(_.jobStatus === GeneratorStatuses.success)
  // 这里直接用arrayLength，第二个参数是数组维度（PostgreSQL数组是1维的，传1）
  .filter(row => arrayLength(row.workerJobIds, 1) === 2)
  .take(N)

如果需要范围判断（比如>0且<N），直接写：

.filter(row => arrayLength(row.workerJobIds, 1) > 0 && arrayLength(row.workerJobIds, 1) < N)

方案2：自定义通用的arrayLength函数（未用slick-pg扩展时） #

如果没引入slick-pg的扩展，需要自定义一个返回Rep[Int]的函数，而不是直接返回布尔值，这样更灵活：

import DatabaseObjects.profile.api._

// 自定义arrayLength函数，返回数组长度的Rep[Int]
def arrayLength[T](arr: Rep[Seq[T]], dimension: Int): Rep[Int] = {
  SimpleExpression.binary[Seq[T], Int, Int] { (arrExpr, dimExpr, queryBuilder) =>
    queryBuilder.sqlBuilder += "array_length("
    queryBuilder.expr(arrExpr)
    queryBuilder.sqlBuilder += ", "
    queryBuilder.expr(dimExpr)
    queryBuilder.sqlBuilder += ")"
  }.apply(arr, LiteralColumn(dimension))
}

val randomFunction: Rep[Double] = SimpleFunction.nullary[Double]("random")
val query: Query[GeneratorV1, GeneratorV1Row, Seq] = GeneratorV1
  .sortBy(_ => randomFunction)
  .filter(_.rasterizationJobConfig +>> "activity" === activityTypeEnum)
  .filter(_.jobStatus === GeneratorStatuses.success)
  .filter(row => arrayLength(row.workerJobIds, 1) === 2)
  .take(N)

为什么之前的代码失败？ #

你之前的自定义函数直接返回了Rep[Boolean]，并且在filter里传了一个元组(row.workerJobIds, arrayLength(row.workerJobIds))，这不符合Slick的filter要求——filter需要接收一个返回Rep[Boolean]的表达式，而不是元组。上面的方案把arrayLength做成返回长度值的表达式，再和数值做比较，就符合Slick的语法要求了。

内容的提问来源于stack exchange，提问作者Matthew Hoggan