Как выбрать строки на основе точного количества элементов массива в другом столбце

#sql #apache-spark #pyspark #apache-spark-sql

#sql #apache-spark #pyspark #apache-spark-sql

Вопрос:

Предположим, у меня есть фрейм данных, подобный этому, где B_C это объединение столбцов B и col C , а столбец selected_B_C — это массив, сформированный путем выбора нескольких B_C столбцов из группы.

  ----------- ----------- -------- -------- ----------------- -------- -------------------------------------- 
|A          |grp_count_A|B       |C       |B_C              |D       |selected_B_C                          |
 ----------- ----------- -------- -------- ----------------- -------- -------------------------------------- 
|1          |6          |30261.41|20091201|30261.41_20091201|99945.83|[30261.41_20091201, 39879.85_20080601]|
|1          |6          |30261.41|20081201|30261.41_20081201|99945.83|[30261.41_20091201, 39879.85_20080601]|
|1          |6          |39879.85|20080601|39879.85_20080601|99945.83|[30261.41_20091201, 39879.85_20080601]|
|1          |6          |69804.42|20080117|69804.42_20080117|99945.83|[30261.41_20091201, 39879.85_20080601]|
|1          |6          |99950.3 |20090301|99950.3_20090301 |99945.83|[30261.41_20091201, 39879.85_20080601]|
|1          |6          |99999.23|20080118|99999.23_20080118|99945.83|[30261.41_20091201, 39879.85_20080601]|
|2          |4          |76498.0 |20150501|76498.0_20150501 |183600.0|[[76498.0_20150501, 76498.0_20150501]]|
|2          |4          |76498.0 |20150501|76498.0_20150501 |183600.0|[[76498.0_20150501, 76498.0_20150501]]|
|2          |4          |76498.0 |20150501|76498.0_20150501 |183600.0|[[76498.0_20150501, 76498.0_20150501]]|
|2          |4          |351378.0|20180620|351378.0_20180620|183600.0|[[76498.0_20150501, 76498.0_20150501]]|
 ----------- ----------- -------- -------- ----------------- -------- -------------------------------------- 
  

Я хочу добавить столбец, selected где он принимает значение 1 , если для строки col B_C находится в col selected_B_C , в противном случае 0 , поэтому конечный фрейм данных выглядит следующим образом.

  ----------- ----------- -------- -------- ----------------- -------- -------------------------------------- -------- 
|A          |grp_count_A|B       |C       |B_C              |D       |selected_B_C                          |selected|
 ----------- ----------- -------- -------- ----------------- -------- -------------------------------------- -------- 
|1          |6          |30261.41|20081201|30261.41_20081201|99945.83|[30261.41_20091201, 39879.85_20080601]|0       |
|1          |6          |30261.41|20091201|30261.41_20091201|99945.83|[30261.41_20091201, 39879.85_20080601]|1       |
|1          |6          |39879.85|20080601|39879.85_20080601|99945.83|[30261.41_20091201, 39879.85_20080601]|1       |
|1          |6          |69804.42|20080117|69804.42_20080117|99945.83|[30261.41_20091201, 39879.85_20080601]|0       |
|1          |6          |99950.3 |20090301|99950.3_20090301 |99945.83|[30261.41_20091201, 39879.85_20080601]|0       |
|1          |6          |99999.23|20080118|99999.23_20080118|99945.83|[30261.41_20091201, 39879.85_20080601]|0       |
|2          |4          |76498.0 |20150501|76498.0_20150501 |183600.0|[[76498.0_20150501, 76498.0_20150501]]|1       |
|2          |4          |76498.0 |20150501|76498.0_20150501 |183600.0|[[76498.0_20150501, 76498.0_20150501]]|1       |
|2          |4          |76498.0 |20150501|76498.0_20150501 |183600.0|[[76498.0_20150501, 76498.0_20150501]]|0       |
|2          |4          |351378.0|20180620|351378.0_20180620|183600.0|[[76498.0_20150501, 76498.0_20150501]]|0       |
 ----------- ----------- -------- -------- ----------------- -------- -------------------------------------- -------- 
  

Сложность для col selected заключается в том, что я хочу, чтобы точное количество вхождений значения в selected_B_C имело значение 1 для selected

  • Например, в group 2 , хотя есть 3 записи со значением 76498.0_20150501 для col B_C , я хочу, чтобы только две записи из group, 2 значение которых 76498.0_20150501 должно иметь значение 1 for selected , поскольку selected_B_C for group 2 имеет ровно 2 элемента со значением 76498.0_20150501 в col selected_B_C

Комментарии:

1. when(col("B_C").isin(flatten(col("selected_B_C")), 1).otherwise(0) Я не тестировал.

2. Спасибо @Lamanus это позволило бы присвоить значение 1 для 9-й записи, но я хочу, чтобы значение для 9-й записи было 0 , поскольку selected_B_C массив для group 2 имеет только два 76498.0_20150501 элемента

3. Я не понимаю значения группы 2. Ах, только первые два из них нужно сделать выбранными!

4. Под группой 2 я подразумеваю группу записей, где значение для столбца A равно 2