Parallel read from db with pyspark

pyspark_parallel_read_from_db.py

import os

q = '(select min(id) as min, max(id) as max from table_name where condition) as bounds'
user = 'postgres'
password = 'secret'
db_driver = 'org.postgresql.Driver'
host = '127.0.0.1'
db_url = f'jdbc:postgresql://{host}:5432/dbname?user={user}&password={password}'
partitions = os.cpu_count() * 2       # a good starting point
conn_properties = {
  'user': 'username',
  'password': 'password',
  'driver': 'org.postgresql.Driver',  # assuming we have Postgres
}

# given that we partition our data by id, get the minimum and the maximum id:
bounds = spark.read.jdbc(
            url=db_url,
            table=q,
            properties=self.conn_properties
        ).collect()[0]

# use the minimum and the maximum id as lowerBound and upperBound and set the numPartitions so that spark
# can parallelize the read from db
df = spark.read.jdbc(
        url=db_url,
        table='(select * from table_name where condition) as table_name',
        numPartitions=partitions,
        column='id',
        lowerBound=bounds.min,
        upperBound=bounds.max + 1,  # upperBound is exclusive
        properties=conn_properties
    )

Hi, I have a concern that as the table is a subquery here, so when JBDC will send this query to database for execution then this subquery will run for every partition and the partition filter created by PySpark will be added later. So, for every parallel run DB will run the complete query as it is a subquery and will add the new filter created by spark later so this activity will be memory intensive at database side.

Please explain if I am not correct

@RK099 interesting 🤔 , thank you for bringing it up, taking a look 👀