How to identify discrete states (oscillations) in Spark Dataframe?

user_6105385 0
Viewed 201

enter image description here

A user U1 moves through the zones Z1, Z2, Z3 at time t1, t2, t3

enter image description here

A user U1 goes back and forth through the zones Z1, Z2 at t1, t2,t3, t4

This is what I call a user « OSCILLATING ». enter image description here

This is considered as an oscillation: the user U1 goes from Z1 to Z2 and then to Z1. The user visits Z1 more than one Time eventhough he visited Z2 only once.

Z1 ==> Z2 ==> Z1

enter image description here

The user U1 goes from Z1 to Z3 then to Z2, Z3 and Z1, respectively at time t1,t2, t3,t4 t5.

The user is oscillation between the 3 zones.

As for the previous example, we consider this movement as an oscillation because the user visits and Z1 and Z3 more than one time eventhough he only visited Z2 only once.

For ease of computation we can set the number of zones max that a user oscillates in to 5 zones.

I would like to give to create a column that tracks the oscillation.

For a given user, if he is oscillating, give the rows the same oscillating ID.

If there is no oscillation, set it to NULL or set it to 0

For example :

enter image description here

Example DATA to copy/paste:

Zone, time, person, Oscillation_ID
A,    1,    ABC,         1
B,    2,    ABC,         1
A,    3,    ABC,         1
A,    4,    ABC,         1
B,    5,    ABC,         1
A,    6,    ABC,         1
C,    7,    ABC,         2
D,    8,    ABC,         2
E,    9,    ABC,         2
C,    10,    ABC,         2
E,    11,    ABC,         2
D,    12,    ABC,         2
C,    13,    ABC,         2
C,    14,    ABC,         2
D,    15,    ABC,         2
E,    16,    ABC,         2
C,    17,    ABC,         2
Z,    18,    ABC,         3
X,    19,    ABC,         4
Y,    20,    ABC,         5

Because I am working with billions of records I would need an efficient solution.

I am using Spark 2.3

I accept both scala and python (pyspark) solutions.

pyspark apache-spark-sql apache-spark sequence
1 Answers

Here's a solution using a window function with Pandas UDF to assign an oscillation id to each row after grouping by person.

I haven't limited max zones in an oscillation as that throws up a bunch of further business logic questions.

I have treated rows to be grouped in the same oscillation until proven otherwise i.e. last two rows in your example dataset are in the same oscillation.

Assuming the input data is ordered by time:

@pandas_udf(IntegerType())
def assign_oscillation(zones: pd.Series) -> int:
  current_oscillation_zones = []
  is_oscillation_frozen = False
  id = 1

  for zone in zones.tolist():
    if zone in current_oscillation_zones and not is_oscillation_frozen:
      is_oscillation_frozen = True
    elif zone not in current_oscillation_zones and is_oscillation_frozen:
      id += 1
      is_oscillation_frozen = False
      current_oscillation_zones = [zone]
    elif zone in current_oscillation_zones and is_oscillation_frozen:
      pass
    elif zone not in current_oscillation_zones and not is_oscillation_frozen:
      current_oscillation_zones.append(zone)
  return id

windowSpec = (Window.partitionBy(col('person'))
              .orderBy(col('time'))
              .rangeBetween(-sys.maxsize, 0))


df.withColumn('Oscillation_ID', assign_oscillation('Zone').over(windowSpec)).show()

I have PySpark 3 and Python 3.8.

PandasUDF within a window function may not be supported in PySpark 2. Here's a less elegant solution using a generator within a PandasUDF with a groupBy that is PySpark 2.3 compatible:

def oscillation_generator(rows: pd.DataFrame) -> pd.DataFrame:
  current_oscillation = pd.DataFrame(data=None)
  current_oscillation_zones = []
  is_oscillation_frozen = False
  id = 1

  for _, row in rows.iterrows():
    if row['Zone'] in current_oscillation_zones[:-1] and not is_oscillation_frozen:
      is_oscillation_frozen = True
      row['Oscillation_ID'] = id
      current_oscillation = current_oscillation.append(row)
    elif row['Zone'] not in current_oscillation_zones and is_oscillation_frozen:
      yield current_oscillation
      id += 1
      is_oscillation_frozen = False
      row['Oscillation_ID'] = id
      current_oscillation = pd.DataFrame(data=[row])
      current_oscillation_zones = [row['Zone']]
    elif row['Zone'] in current_oscillation_zones and is_oscillation_frozen:
      row['Oscillation_ID'] = id
      current_oscillation = current_oscillation.append(row)
    elif row[
      'Zone'] not in current_oscillation_zones and not is_oscillation_frozen:
      current_oscillation_zones.append(row['Zone'])
      row['Oscillation_ID'] = id
      current_oscillation = current_oscillation.append(row)
  yield current_oscillation

@pandas_udf(StructType(
  [StructField('Zone', StringType()),
   StructField('time', IntegerType()),
   StructField('person', StringType()),
   StructField('Oscillation_ID', IntegerType())]), PandasUDFType.GROUPED_MAP)
def assign_oscillations(rows: pd.DataFrame) -> pd.DataFrame:
  oscillations = oscillation_generator(rows)
  return pd.concat([oscillation for oscillation in oscillations])


df.groupBy(['person']).apply(assign_oscillations).show()
user_18593700
Related