ENH: option efficient memory use by downcasting

[dkapitan]

We use pandas-gbq a lot for our daily analyses. It is known that memory consumption can be a pain, see e.g. https://www.dataquest.io/blog/pandas-big-data/

I have started to write a patch, which could be integrated into an enhancement for read_gbq (rough idea, details TBD):

• Provide boolean optimize_memory option
• If True, the source table is inspected with a query to get min, max, presence of nulls and % of unique number of strings for INTEGER and STRING columns, respectively
• When calling to_dataframe this information is passed to the dtypes option, downcasting integers to the appropriate numpy (u)int type, and converting strings to pandas category type at some threshold (less than 50% of unique values)

I already have a working monkey-patch, which is still a bit rough. If there is enough interest I'd happily make it more robust and submit a PR. Would be my first significant contribution to an open source project, so some help and feedback would be appreciated.

Curious to hear your views on this.

[max-sixty]

Thanks for the suggestion @dkapitan

Interesting ideas - what's the rationale for not doing that every time? I haven't used the new int dtype much; my understanding is that categoricals were generally dominant of numpy object strings, it's basically dictionary encoding

Either way, I would definitely be keen to add these as options. Explicit > implicit, and I think that having the library behind the scenes send additional queries and self-optimize would be too much magic, at least to start. What are your thoughts?

[dkapitan]

@max-sixty
Thanks for the prompt response, motivates me to get this going not just for ourselves. Agree that this feature could be used more generically, and that for the first iteration we should explicitly separate the steps (rather than opaque magic). So here's my user-story:

Enhancements

• GbqConnector.optimize_schema(table_id, threshold_strings=0.5, default_float=np.float32):

  1. reflects on existing table or query results to determine current SQL types
  2. does calculation in standardSQL to determine optimal pandas dtypes
  3. returns dtypes as dict, possibly with table with statistics from ii) for inspection

• Add read_gbq(dtypes=None) option, for passing result of iii) to google.cloud.bigquery.job.QueryJob.to_dataframe(dtypes=dtypes)

Design considerations and focus first version

• optimze_schema only on existing tables, reflection of dynamic query results requires more engineering (e.g. referencing temporary table in GBQ cache)
• Only optimize non-nested, non-repeated columns. Pandas can’t handle these anyhow, so any nested and/or repeated columns will be read as a dict in an object column
• Logic for downcasting INTEGER and STRING. Need to make trade-off between null-handling of INTEGER since only pandas.Int64 is nullable, numpy ints aren't. This is actually quite significant: lot's of coded.enumerated data often fits in uint8 or uint16 but then we need to see how we handle nulls (use sentinel values?!)
• FLOAT can be set to default np.float32 without inspection. Almost for all my use cases I don't need larger mantissa or exponent.
• User need to explicitly invoke optimize_schema and pass it in dtypes. In future we may want to consider to_dataframe(optimize=True)

Feel free to add/modify so we have a clear goal for the first iteration. Have started with the how-to for contributors, my fork is at https://github.com/dkapitan/pandas-gbq. Will let you know when I have done my first commit of the optimize-memory-usage branch.

[tswast]

Explicit > implicit, and I think that having the library behind the scenes send additional queries and self-optimize would be too much magic, at least to start.

I agree. Especially since it's not supported by the BigQuery team to query anonymous query results tables.

I propose we add a dtypes argument to pandas-gbq (passing through to that of google-cloud-bigquery) to support this use case. Maybe we make a public function to get a memory optimal schema via these queries that can be used to pass as the dtypes argument?

[tswast]

Oops, I should read more closely, I think I just proposed your user story. :-)

[max-sixty]

Thanks for the clear plan @dkapitan (I've just returned from a conference so thanks in advance for your patience)

Can we clarify whether the data in the data influences whether we'd want these - e.g. why not always use pandas.Int64, or categorical for strings?

If this is something that doesn't depend on the underlying data, we could skip the work behind GbqConnector.optimize_schema. We could add the dtypes option first, and then potentially promote to default after more experience.

[tswast]

The blog post recommends using smaller types where possible.

The way this is implemented in google-cloud-bigquery, I fear you'll still need quite a bit of temporary memory (that should get freed after dataframe construction), since it does convert the BQ response JSON/protos into Python objects first. I'm working on improving this situation with the BQ Storage API via a custom Avro parser (googleapis/google-cloud-python#7805), but even then we'll use the most natural dtype first (e.g. pandas int64 for BigQuery INTEGER columns). Using a smaller dtype will require a copy operation which will slow the parsing down, but will save steady-state memory (but still require enough temporary memory to fit the column as int64). I guess the Avro parser could use the requested dtype to avoid this, but then you risk either out-of-bounds errors or silent overflow.

[dkapitan]

@tswast @max-sixty

So the good news is that I have got a working version, see this commit. Basic testing seems OK (haven't written all the tests yet), but when actually converting the query results to a dataframe I'm getting weird results. It seems that, depending on the size of the QueryJob.results, the conversion from STRING to category fails; it only works for small result sets. See attached .ipynb. I have tried to trace this bug in the source code, but can't find anything. Your help would be much appreciated.

Then, regarding your questions:

• Donwncasting INTEGER to smaller numpy.int types significantly reduced the memory footprint
• Indeed during load-time there will be a temporary peak in memory usage. It would be best if the type conversion is done directly from Avro to pandas, but indeed you risk errors as mentioned. I've actually come across that during testing: integer data got all messed up with negative numbers because I'd enforced an int-type that was too small.

Finally, on a practical note: I don't have much experience using Git with branches (only used basic Git stuff). I'm afraid I've made a mess of my branch already (too many merges and rebases whilst trying to figure out how it works). Please let me know if this OK (assuming it can be fixed during a PR). If not, please let me know and I will start with a clean fork.

Curious to hear your thoughts.

pandas-gbq-bugtracing.ipynb.zip

[tswast]

the conversion from STRING to category fails; it only works for small result sets.

Probably because the downcasting happens for each page, so if there's a page without all the possible categories, the concat will fail.

[tswast]

Please let me know if this OK (assuming it can be fixed during a PR).

I can probably clean it up if you mail the PR and check the "allow edits from maintainers" box. Don't worry about it too much.

[tswast]

the conversion from STRING to category fails; it only works for small result sets.

Probably because the downcasting happens for each page, so if there's a page without all the possible categories, the concat will fail.

Could you file an issue to google-cloud-python with this for me? Maybe there's a way we can delay the cast. I hadn't actually considered the categorical type when I added the dtype option.

[max-sixty]

Great @dkapitan

My experience is that pandas doesn't deal that well with non-standard types (e.g. smaller floats are upcast on some operations). Combined with our current parsing, I'm not sure it'll save much memory, consistent with @tswast 's thoughts above.

If there's any tradeoff in prioritization between @dkapitan suggestions around using 1) nullable ints and 2) categoricals for strings; those would be immediately impactful.

[max-sixty]

Maybe there's a way we can delay the cast.

💯

[dkapitan]

Thanks for all the feedback @max-sixty @tswast.

I am not sure how to proceed from here. Below some extra tests, thoughts and questions. Please let me know what you think is best to do next.

Persistence of downcasted types

My experience is that pandas doesn't deal that well with non-standard types (e.g. smaller floats are upcast on some operations).

Have done a simple test (see below), and for now I'd say we stick to the following:

• if we use np.float32 throughout, upcasting by pandas should be reduced since all float columns have been coerced to np.float32. Of course there may be less frequently used methods in pandas which would still upcast to float64, but I expect this would reduce memory usage in most use cases
• for now I'd prefer to stick to pandas nullable Int64 for integer columns containing null-values, at the cost of more memory. I expect users requiring more memory reduction would do further optimisation themselves, replacing nulls with sentinel values and downcast from there. Integer columns with no nulls are downcasted to the smallest size

# test upcasting by pandas
df = (pd.DataFrame({'_float64': [1.1, 2.2, None, 4.4],
                    '_float32': [1.1, None, 3.3, 4.4],
                    '_float16': [1.1, 2.2, 3.3, 4.4],
                    '_int64': [1, None, 3, 4],
                    '_int32': [1, 2, -3, 4],
                    '_uint16': [1, 2, 3, 4],
                    })
      .astype({'_float64': 'float',
               '_float32': np.float32,
               '_float16': np.float16,
               '_int64': 'Int64',
               '_int32': np.int32,
               '_uint16': np.uint16,
               })
      )
df['_sum_floats_16'] = df._float16 + df._float16
df['_sum_floats_32'] = df._float32 + df._float32
df['_sum_floats_mixed'] = df._float16 + df._float32
df['_sum_ints_mixed'] = df._int32 + df._uint16
print(pd.concat([df.dtypes, df.memory_usage()], axis=1, sort=False)
      .rename(columns=({0: 'dtype', 1: 'mem_use'})))

# output
                     dtype  mem_use
_float64           float64       32
_float32           float32       16
_float16           float16        8
_int64               Int64       36
_int32               int32       16
_uint16             uint16        8
_sum_floats_16     float16        8
_sum_floats_32     float32       16
_sum_floats_mixed  float32       16
_sum_ints_mixed      int32       16
Index                  NaN       80

Peak memory consumption
I have found the memory-profiler module to plot the memory usage over time. Executing mprof mem_test.py yields the following graph, showing peak memory use is 3.5 x compared to post-load.

memory usage with optimized dataframe (df_optimized 0.8 MB)

memory usage with non-optimized dataframe (df 4.6 MB)

Although df_optimized is reduced to 17% of the original df size for this (small) dataset, this reduction is negligible on the total memory footprint of the script. The impact may be larger for large datasets, but then it defeats the whole purpose of this exercise since the peak memory usage is likely to exceed the systems RAM (assuming something like 16 GB RAM for a data science machine/laptop).

#!/usr/bin/env python
# coding: utf-8
# mem_test.py

from pandas_gbq.gbq import read_gbq, GbqConnector, optimize_dtypes
from pandas_gbq.schema import (
    select_columns_by_type,
    generate_sql,
    _determine_int_type,
    _determine_string_type,
)
import pandas as pd
import numpy as np


default_project_id = "mediquest-cloud"


def count_dtypes(df):
    return df.dtypes.apply(lambda x: x.name).value_counts()


def run_mem_test():
    for table in [
        "baseball.games_post_wide",
        #              'ml_datasets.credit_card_default',
    ]:
        table_reference = ".".join(["bigquery-public-data", table])
        optimized = optimize_dtypes(
            table_reference, default_project_id=default_project_id
        )
    #     df_original = read_gbq(f'SELECT * FROM `{table_reference}`', use_bqstorage_api=True)
    #     df_optimized = read_gbq(f'SELECT * FROM `{table_reference}`', use_bqstorage_api=True, dtypes=optimized['dtypes'])
    #     print('Memory usage optimized dataframe: {:.1f}% of orginal dataframe'.
    #           format(100*df_optimized.memory_usage().sum()/df_original.memory_usage().sum()))

    # strange behaviour: string columns are converted to category when 'limit 4000' in query
    sql = "select * from `bigquery-public-data.baseball.games_post_wide` limit {}"
    print("Reading from GBQ: {}".format(table))
    df1 = read_gbq(sql.format(4000), dtypes=optimized["dtypes"])
    print(count_dtypes(df1))
    print(
        "Memory use optimized table {}: {}".format(
            table, df1.memory_usage().sum()
        )
    )


if __name__ == "__main__":
    run_mem_test()

STRING -> category conversion
I will await the solution for BigQuery GH8044 for conversion of strings to categories