在pyspark列中查找列表中连续的长度

1 回答

• 1

  Spark版本2.1及以上版本

  如果你有Spark版本2.1或更高版本，这是一种方法：

  首先使用pyspark.sql.posexplode()将收集的列表与索引一起展开 .

  import pyspark.sql.functions as f
  df = df.select("a", f.posexplode("b").alias("pos", "b"))
  

  接下来使用Window函数创建一个列，该列将指示当前行的值是否与前一行不同 .

  from pyspark.sql import Window
  
  w = Window.partitionBy("a").orderBy("pos")
  df = df.select(
      "*", 
      (f.col("b") != f.lag(f.col("b"), default=0).over(w)).cast("int").alias("change")
  )
  df.show()
  #+----+---+---+------+
  #|   a|pos|  b|change|
  #+----+---+---+------+
  #|bank|  0|  1|     1|
  #|bank|  1|  0|     1|
  #|bank|  2|  1|     1|
  #|bank|  3|  1|     0|
  #|bank|  4|  1|     0|
  #|bank|  5|  1|     0|
  #|tank|  0|  1|     1|
  #|tank|  1|  1|     0|
  #|tank|  2|  0|     1|
  #|tank|  3|  1|     1|
  #|tank|  4|  1|     0|
  #|tank|  5|  1|     0|
  #|tank|  6|  1|     0|
  #+----+---+---+------+
  

  计算该列的累积总和，将 0 和 1 分成组 . 然后你可以 groupBy() 列 (a, b, group) 并计算每个组的长度以及开始和结束索引 .

  df = df.select(
      "*",
      f.sum(f.col("change")).over(w.rangeBetween(Window.unboundedPreceding, 0)).alias("group")
  )\
  .groupBy("a", "b", "group")\
  .agg(f.min("pos").alias("start"), f.max("pos").alias("end"), f.count("*").alias("length"))\
  .where(f.col("b") == 1)\
  .drop("group")
  df.show()
  #+----+---+-----+---+------+
  #|   a|  b|start|end|length|
  #+----+---+-----+---+------+
  #|bank|  1|    0|  0|     1|
  #|bank|  1|    2|  5|     4|
  #|tank|  1|    0|  1|     2|
  #|tank|  1|    3|  6|     4|
  #+----+---+-----+---+------+
  

  最后，您可以过滤此DataFrame以查找与 a 列中每个项目的最长长度序列关联的行：

  df = df.withColumn(
      "isMax",
      f.col("length") == f.max(f.col("length")).over(Window.partitionBy("a"))
  )\
  .where(f.col("isMax"))\
  .drop("isMax")
  df.show()
  #+----+---+-----+---+------+
  #|   a|  b|start|end|length|
  #+----+---+-----+---+------+
  #|bank|  1|    2|  5|     4|
  #|tank|  1|    3|  6|     4|
  #+----+---+-----+---+------+
  

  ---

  Spark版本1.5及以上版本

  如果你没有 posexplode ，另一种选择是将你的整数数组转换成一个字符串数组，连接它，并在 "0" 上拆分 . 然后爆炸结果数组，并筛选具有最大长度的数组 .

  不幸的是，这种方法并没有给你起点和终点 .

  df.withColumn('b', f.split(f.concat_ws('', f.col('b').cast('array<string>')), '0'))\
      .select('a', f.explode('b').alias('b'))\
      .select('a', f.length('b').alias('length'))\
      .withColumn(
          "isMax",
          f.col('length') == f.max(f.col('length')).over(Window.partitionBy('a'))
      )\
      .where(f.col("isMax"))\
      .drop("isMax")\
      .show()
  #+----+------+
  #|   a|length|
  #+----+------+
  #|bank|     4|
  #|tank|     4|
  #+----+------+
  

  回复于 2024-05-06T23:50:51+08:00