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.

💯