PySpark实现Aroon指标：按组计算最值与当前行的间隔行数

Got it, let's build out the missing pieces to calculate those period gaps and finish the Aroon indicators. Here's how to extend your existing PySpark code step by step:

First, we need to track row positions within each group so we can calculate the distance between the current row and the row where the max/min occurred. Then we'll compute the period gaps, and finally derive the Aroon metrics.

Step 1: Import Required Functions #

Make sure you have these imports at the top of your script:

from pyspark.sql.functions import row_number, first, when, max, min
from pyspark.sql.window import Window

Step 2: Add Grouped Row Numbers #

We'll assign a sequential row number to each entry within its Name group, ordered by Date—this helps us measure the row interval:

def myFunction(myData):
    df = myData
    col = 'Date'
    groupCols = ['Name']
    
    # Add row number within each group (ordered by Date)
    row_window = Window.partitionBy(groupCols).orderBy(df[col].asc())
    df = df.withColumn("row_num", row_number().over(row_window))
    
    # Your existing window and max/min calculations
    window=Window.partitionBy(groupCols).orderBy(df[col].asc()).rowsBetween(-11, 0)
    max_value = max(df['value']).over(window)
    min_value = min(df['value']).over(window)
    df = df.withColumn('max', max_value)
    df = df.withColumn('min', min_value)

Step 3: Calculate Periods Since Max/Min #

Next, we'll find the row number of the most recent max/min in the sliding window, then compute the interval between that row and the current row:

# Get row number of the most recent max value in the window
    max_row_window = window.orderBy(when(df["value"] == df["max"], 0).otherwise(1), df["row_num"].desc())
    max_row_num = first(df["row_num"], ignorenulls=True).over(max_row_window)
    df = df.withColumn("periods_since_max", df["row_num"] - max_row_num)
    
    # Get row number of the most recent min value in the window
    min_row_window = window.orderBy(when(df["value"] == df["min"], 0).otherwise(1), df["row_num"].desc())
    min_row_num = first(df["row_num"], ignorenulls=True).over(min_row_window)
    df = df.withColumn("periods_since_min", df["row_num"] - min_row_num)

Note: The orderBy clause ensures we prioritize rows where the value matches the window's max/min, and if there are duplicates, we pick the most recent one (highest row number).

Step 4: Compute Aroon Indicators #

Now we can calculate the final Aroon metrics using the period gaps:

# Calculate Aroon Up, Aroon Dn, and Aroon Oscillator
    df = df.withColumn("aroon_up", ((12 - df["periods_since_max"]) / 12) * 100)
    df = df.withColumn("aroon_dn", ((12 - df["periods_since_min"]) / 12) * 100)
    df = df.withColumn("aroon_oscillator", df["aroon_up"] - df["aroon_dn"])
    
    return df

Full Combined Code #

Putting it all together, the complete function looks like this:

from pyspark.sql.functions import row_number, first, when, max, min
from pyspark.sql.window import Window

def myFunction(myData):
    df = myData
    col = 'Date'
    groupCols = ['Name']
    
    # Add row number within each group (ordered by Date)
    row_window = Window.partitionBy(groupCols).orderBy(df[col].asc())
    df = df.withColumn("row_num", row_number().over(row_window))
    
    # Sliding window for max/min calculations (12-period window: current + previous 11 rows)
    window=Window.partitionBy(groupCols).orderBy(df[col].asc()).rowsBetween(-11, 0)
    max_value = max(df['value']).over(window)
    min_value = min(df['value']).over(window)
    df = df.withColumn('max', max_value)
    df = df.withColumn('min', min_value)
    
    # Get row number of most recent max in window
    max_row_window = window.orderBy(when(df["value"] == df["max"], 0).otherwise(1), df["row_num"].desc())
    max_row_num = first(df["row_num"], ignorenulls=True).over(max_row_window)
    df = df.withColumn("periods_since_max", df["row_num"] - max_row_num)
    
    # Get row number of most recent min in window
    min_row_window = window.orderBy(when(df["value"] == df["min"], 0).otherwise(1), df["row_num"].desc())
    min_row_num = first(df["row_num"], ignorenulls=True).over(min_row_window)
    df = df.withColumn("periods_since_min", df["row_num"] - min_row_num)
    
    # Calculate Aroon indicators
    df = df.withColumn("aroon_up", ((12 - df["periods_since_max"]) / 12) * 100)
    df = df.withColumn("aroon_dn", ((12 - df["periods_since_min"]) / 12) * 100)
    df = df.withColumn("aroon_oscillator", df["aroon_up"] - df["aroon_dn"])
    
    return df

How It Works #

• Row Numbers: We use row_number() to assign a unique sequence to each row in its group, ordered by date—this lets us measure how many rows have passed since the max/min.
• Finding Recent Max/Min: The custom window ordering ensures we grab the most recent occurrence of the window's max/min value (critical if there are duplicate peaks/troughs).
• Aroon Calculations: The formulas follow the standard Aroon definitions:
  • Aroon Up measures how long ago the highest high occurred in the window (100 means it's the current row).
  • Aroon Dn does the same for the lowest low.
  • Aroon Oscillator is the difference between the two, showing trend strength and direction.

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