Using IN Operator or isin Function

Let us understand how to use IN operator while filtering data using a column against multiple values.

  • It is alternative for Boolean OR where single column is compared with multiple values using equal condition.

Let us start spark context for this Notebook so that we can execute the code provided. You can sign up for our 10 node state of the art cluster/labs to learn Spark SQL using our unique integrated LMS.

from pyspark.sql import SparkSession

import getpass
username = getpass.getuser()

spark = SparkSession. \
    builder. \
    config('spark.ui.port', '0'). \
    config("spark.sql.warehouse.dir", f"/user/{username}/warehouse"). \
    enableHiveSupport(). \
    appName(f'{username} | Python - Basic Transformations'). \
    master('yarn'). \
    getOrCreate()

If you are going to use CLIs, you can use Spark SQL using one of the 3 approaches.

Using Spark SQL

spark2-sql \
    --master yarn \
    --conf spark.ui.port=0 \
    --conf spark.sql.warehouse.dir=/user/${USER}/warehouse

Using Scala

spark2-shell \
    --master yarn \
    --conf spark.ui.port=0 \
    --conf spark.sql.warehouse.dir=/user/${USER}/warehouse

Using Pyspark

pyspark2 \
    --master yarn \
    --conf spark.ui.port=0 \
    --conf spark.sql.warehouse.dir=/user/${USER}/warehouse

Tasks

Let us perform some tasks to understand filtering in detail. Solve all the problems by passing conditions using both SQL Style as well as API Style.

  • Read the data for the month of 2008 January.

airtraffic_path = "/public/airtraffic_all/airtraffic-part/flightmonth=200801"
airtraffic = spark. \
    read. \
    parquet(airtraffic_path)
airtraffic.printSchema()
root
 |-- Year: integer (nullable = true)
 |-- Month: integer (nullable = true)
 |-- DayofMonth: integer (nullable = true)
 |-- DayOfWeek: integer (nullable = true)
 |-- DepTime: string (nullable = true)
 |-- CRSDepTime: integer (nullable = true)
 |-- ArrTime: string (nullable = true)
 |-- CRSArrTime: integer (nullable = true)
 |-- UniqueCarrier: string (nullable = true)
 |-- FlightNum: integer (nullable = true)
 |-- TailNum: string (nullable = true)
 |-- ActualElapsedTime: string (nullable = true)
 |-- CRSElapsedTime: integer (nullable = true)
 |-- AirTime: string (nullable = true)
 |-- ArrDelay: string (nullable = true)
 |-- DepDelay: string (nullable = true)
 |-- Origin: string (nullable = true)
 |-- Dest: string (nullable = true)
 |-- Distance: string (nullable = true)
 |-- TaxiIn: string (nullable = true)
 |-- TaxiOut: string (nullable = true)
 |-- Cancelled: integer (nullable = true)
 |-- CancellationCode: string (nullable = true)
 |-- Diverted: integer (nullable = true)
 |-- CarrierDelay: string (nullable = true)
 |-- WeatherDelay: string (nullable = true)
 |-- NASDelay: string (nullable = true)
 |-- SecurityDelay: string (nullable = true)
 |-- LateAircraftDelay: string (nullable = true)
 |-- IsArrDelayed: string (nullable = true)
 |-- IsDepDelayed: string (nullable = true)
  • Get count of flights departed from following major airports - ORD, DFW, ATL, LAX, SFO.

airtraffic. \
    filter("Origin IN ('ORD', 'DFW', 'ATL', 'LAX', 'SFO')"). \
    count()
118212
airtraffic.count()
605659
  • API Style

from pyspark.sql.functions import col
c = col('x')
help(c.isin)
Help on method isin in module pyspark.sql.column:

isin(*cols) method of pyspark.sql.column.Column instance
    A boolean expression that is evaluated to true if the value of this
    expression is contained by the evaluated values of the arguments.
    
    >>> df[df.name.isin("Bob", "Mike")].collect()
    [Row(age=5, name='Bob')]
    >>> df[df.age.isin([1, 2, 3])].collect()
    [Row(age=2, name='Alice')]
    
    .. versionadded:: 1.5
from pyspark.sql.functions import col

airtraffic. \
    filter(col("Origin").isin("ORD", "DFW", "ATL", "LAX", "SFO")). \
    count()
118212
  • Get number of flights departed late on Sundays as well as on Saturdays. We can solve such kind of problems using IN operator as well.

from pyspark.sql.functions import col, concat, lpad

airtraffic. \
    withColumn("FlightDate",
               concat(col("Year"),
                      lpad(col("Month"), 2, "0"),
                      lpad(col("DayOfMonth"), 2, "0")
                     )
              ). \
    filter("""
           IsDepDelayed = 'YES' AND Cancelled = 0 AND
           date_format(to_date(FlightDate, 'yyyyMMdd'), 'EEEE') IN
               ('Saturday', 'Sunday')
           """). \
    count()
57873
  • API Style

from pyspark.sql.functions import col, concat, lpad, date_format, to_date

airtraffic. \
    withColumn("FlightDate",
               concat(col("Year"),
                      lpad(col("Month"), 2, "0"),
                      lpad(col("DayOfMonth"), 2, "0")
                     )
              ). \
    filter((col("IsDepDelayed") == "YES") & (col("Cancelled") == 0) &
           (date_format(
               to_date("FlightDate", "yyyyMMdd"), "EEEE"
           ).isin("Saturday", "Sunday"))
          ). \
    count()
57873