Herein, I will outline the several “gotchas” I have stumbled upon in my quest to read data from an SQL table using python pandas in a memory efficient way. In my case, the table was 52M rows, 23 columns, and used 900 MB RAM when loaded, but uses north of 12 GB when fetching it from SQL, which was more than the production environment ideally can handle. In this post I detail different solutions that each will reduce memory consumption.

Below is the first generic starting code which I want to improve. All other snippets are additions/modifications of the below code. (It is assumed you have a working connection to SQL from python from before. Here, we use luigi to read in needed parameters.)

import pandas as pd
import pyodbc

# We use luigi to get the SQL user, password, database, and driver.
# It's assumed you can conjure up correct information to establish a
# pyodbc connection to SQL
import luigi

def get_db():
    """Returns connection object to SQL"""
    config = luigi.configuration.get_config()
    kwargs = dict(config.items("database"))
    return pyodbc.connect(**kwargs)

# Table name to be read
table = "my_table"

# Get connection to SQL database
conn = get_db()

# Select only the columns we actually need, from SQL
columns = ["COLUMN01",
           "COLUMN02",
           # ....
           "COLUMN23"]

query = f"""SELECT {', '.join(columns)} FROM {table}"""

df = pd.read_sql(query, conn).reset_index(drop=True)

Since the end goal is to run this in a docker image in production environment, we do all bench marking, such as RAM usage, by monitoring the running docker image, as described in a previous post. We see it reaches 12 GB of RAM use in the figure below (we sample the RAM usage once per second). For interest, we also include CPU usage.

After all the data is loaded from SQL, it is validated using Great Expectations, which is the second half of activity seen below. When this is done, it returns, and the docker image stops, which here takes under 2 minutes.

Pandas’ read_sql() has an argument chunksize, if given it will return the data in chunks. We will then have to concatenate each df, corresponding to the chunks.

Since we are appending, we need to initialize a data frame, and specify the columns.

df = pd.DataFrame(columns=columns)

chunksize = 100_000
for i, df_chunk in enumerate(pd.read_sql(query, conn, chunksize=chunksize)):
    df = pd.concat((df, df_chunk), axis=0, copy=False)
    print(f"{i} Got dataframe with {len(df_chunk):,} rows")

df = df.reset_index(drop=True)

With this modification, RAM usage goes down and now peaks at 8 GB. But this is not one 50th of what we had when we read in all 52M rows at once, why?

It turns out, when setting chunksize, it just returns it in chunks, it still reads and loads the table all at once, as detailed in this blog post.

This is fixed by using sqlalchemy, instead of pyodbc to establish a connection, and setting stream_results=True (there is an open issue on pandas to do this by default).

from sqlalchemy import create_engine
import luigi

def get_db():
    """Returns connection object to SQL"""
    config = luigi.configuration.get_config()
    kwargs = dict(config.items("database"))

    engine_stmt = ("mssql+pyodbc://%s:%s@%s/%s?driver=%s" % \
                   (kwargs['uid'],
                    kwargs['pwd'],
                    kwargs['server'],
                    kwargs['database'],
                    kwargs['driver'])) # 'ODBC+DRIVER+17+for+SQL+Server'

    # Query in streaming chunk size, to save RAM.
    engine = create_engine(engine_stmt)
    conn = engine.connect().execution_options(stream_results=True)
    return conn

However, this did not get me below 6GB of RAM usage (results not shown).

Depending on the data, significant RAM can be saved by casting each column into more suitable data types, especially if few values are repeated often, such that we can cast to category type. However, since we are concatenating the chunked reading, the “top”/original df and the “bottom”/appending one need to be in agreement with which categories map, internally, to which placeholder index, which we now address here.

It turns out that the real offending culprit is pd.concat(), which uses multiple times more RAM than the resulting df. The workaround is to not use it, e.g. by concatenating manually to disk, instead of to RAM, as pd.concat does. This is addressed in this post where they use read_csv(), but since we are reading from sql, we use the code snippet below.

Note: when incrementally fetching data with an OFFSET, it is important the column we are telling SQL to do its ordering on is all unique values, or there will be repeating rows, from the different chunks.

import pyarrow
import pyarrow.parquet as pq

chunksize, chunk = 100_000, 0

# Note: values in ordering column (here: 'COLUMN01') must be all unique:
q = query + " ORDER BY COLUMN01 OFFSET {} ROWS FETCH NEXT {} ROWS ONLY"

while True:
    df_tmp = pd.read_sql_query(q.format(chunksize*chunk, chunksize), conn)

    # Save df to parquet by appending
    table = pyarrow.Table.from_pandas(df_tmp)
    pq.write_to_dataset(table, root_path="disk_data", use_deprecated_int96_timestamps=True)

    print(f"Got chunk {chunk} of size {len(df_tmp):,}")

    if len(df_tmp) < chunksize:
        break
    chunk += 1

df = pd.read_parquet("disk_data").reset_index(drop=True)

Note, "disk_data" here is the name of a folder. One can wrap the block in a with tempfile.TemporaryDirectory() as disk_data block.

This used 500 MiB RAM during reading from SQL, and then a spike just below 4 GB when reading the data back, and validating it using Great Expectations, as mentioned in the 1.

Comparing all methods together, its more clear that there is also a trade off with speed. Appending to disk is the slowest, but saves the most memory usage. The increase in RAM towards the end is due to Great Expectations reading in the full data and running validations on it.