Создать новый столбец с групповой меткой на основе условия, проверенного в другом столбце

Question

Я хочу создать групповые метки на основе условия, проверенного в другом столбце.В частности, если есть значение directionChange.equalTo(1), я хочу начать новый сегмент (метку).Результат должен быть указан в столбце segmentNr результат, полученный моим кодом, находится в nSegment.Я предполагаю, что это невозможно для подобных заданий.Наконец, я хотел бы рассчитать различные аспекты, такие как сумма, среднее, максимальное значение на сегмент (не входит в объем вопроса).

Пример ввода:

+---+-----+---------------+---------+--------+
| id|value|directionChange|segmentNr|nSegment|
+---+-----+---------------+---------+--------+
|  1| 11.0|              0|        1|       1|
|  2|-22.0|              1|        2|       1|
|  3| 34.0|              0|        2|       1|    
|  4|-47.0|              1|        3|       1|    
|  5| 61.0|              1|        4|       1|    
|  6| 22.0|              0|        4|       1|    
|  7|  5.0|              0|        4|       1|    
|  8| -7.0|              1|        5|       1|    
+---+-----+---------------+---------+--------+

Функция добавления нового столбца с использованием входного набора данных:

public static Dataset<Row> createSegments(Dataset<Row> dataset, String columnName, int start, String newColumnName) throws Exception
{
    int test = 1;
    Dataset<Row> resultDataset = dataset.withColumn(newColumnName, //
            functions.when(dataset.col(columnName).equalTo(1), (start = start + 1))//
                    .otherwise(start));

    return resultDataset;
}

Функция вызывается следующим образом:

dataset = createSegments(dataset, "directionChange", 0, "nSegment");

Answer 1

Это можно сделать с помощью функции Window. Однако, поскольку у вас нет столбца для разделения данных, он может стать очень медленным для больших наборов данных. Это можно улучшить, используя partitionBy(column) для объекта Window ниже. Однако для этого потребуется хороший столбец, и окончательный результат также будет разбит на части.

Идея решения состоит в том, чтобы сделать кумулятивную сумму столбца directionChange при заказе по столбцу id. В Scala:

val window = Window.orderBy("id").rowsBetween(Window.unboundedPreceding, Window.currentRow)
val df2 = dataset.withColumn("nSegment", sum($"directionChange").over(window) + 1)

Java-код:

WindowSpec window = Window.orderBy("id").rowsBetween(Window.unboundedPreceding(), Window.currentRow()); 
Dataset<Row> df2 = dataset.withColumn("nSegment", functions.sum("directionChange").over(window));

В старых версиях Spark (<2.1.0) используйте: </p>

rowsBetween(Long.MinValue, 0)

Это создаст новый столбец nSegment, равный segmentNr из входных данных. Для следующего шага вы можете использовать groupBy("nSegment").agg(...) для расчета различных показателей для каждого сегмента.