Как выбрать строки по условию общего значения указанного столбца c?

Question

Извините за запутанный заголовок, я не знаю, как его выразить.

Это то, что я хочу сделать, sh с помощью PySpark SQL:

Когда Значение для переменной 'Z' составляет от 2000 до 3000 строк, возвращаемых для указанного c ID, но возвращаются только данные для переменных Y и Z.

Я не знаю, как go дальше, чем с запрос ниже, как сообщить SQL, что нам нужно выбрать правильные идентификаторы (2 и 3) в прикрепленном df?

 SELECT ID, Variable, Date, Value 
 FROM TABLE 
 WHERE (Variable == 'Y' OR Variable == 'Z') AND "if Value of Z between 2000 and 3000 then select only these IDs"

Вверху - df, внизу - ожидаемый результат

Answer 1

В Spark я бы рекомендовал использовать оконные функции:

select t.*
from (select t.*,
             sum(case when variable = 'Z' and z_value between 2000 and 3000 then 1 else 0 end) over (partition by id) as z_value_cnt
      from t
     ) t
where variable in ('Y', 'Z') and z_value_cnt >= 1;

Я ожидал бы, что это будет иметь лучший план выполнения.

Answer 2

Попробуйте использовать простой внутренний запрос

 SELECT ID, Variable, Date, Value 
 FROM TABLE 
 WHERE (Variable == 'Y' OR Variable == 'Z') AND ID in (select id from TABLE where variable = 'Z' and value between 2000 and 3000)

Answer 3

df=spark.createDataFrame([(1,'X', '01/01/20',5),
                      (1,'Y', '01/01/20',10),
                      (1,'Z', '01/01/20',1000),
                      (2,'X', '01/02/20',20),
                      (2,'Y', '01/02/20',30),
                      (2,'Z', '01/02/20',2500)],['ID','VAR','DATE','VAL'])
display(df.filter("VAL Between 2000 and 3000").select('ID').join(df,['ID']).filter("VAR in ('X','Y')"))