Optimize postgres performance #531
Conversation
dimberman
force-pushed
the
postgres-imp-2
branch
from
8a6f337
to
f91bbd7
Compare
| @@ -63,7 +63,7 @@ def create_dag(database_name, table_args, dataset): | |||
| output_table=table, | |||
| chunk_size=chunk_size, | |||
| ) | |||
| aql.truncate(my_table) | |||
| aql.cleanup([my_table]) | |||
There was a problem hiding this comment.
Why have we introduced this change?
There was a problem hiding this comment.
@tatiana two reasons
- there are issues in truncate that I'm gonna make tickets for today
- Truncate empties but doesn't delete the table. We want to delete the table here.
There was a problem hiding this comment.
@dimberman Thank you! :)
I see, since we are not naming these tables, that will work.
Moving forward, we may have to add an aql.drop since cleanup does not drop tables which are not temporary. What do you think? Should we log a ticket for this? Another possibility would be to have a parameter in cleanup also to delete non-temporary tables..!
There was a problem hiding this comment.
BTW: are there more bugs besides #515?
There was a problem hiding this comment.
@tatiana that's the bug I was referring to though I don't think it's limited to just load_file.
It's a pretty simple fix too. Just need to set template_fields = ("table",) in the TruncateOperator class. I can do that this afternoon or tomorrow if no one else gets around to it.
| ], | ||
| indirect=True, | ||
| ids=["snowflake", "bigquery", "postgresql"], | ||
| ids=["snowflake", "bigquery"], | ||
| ) | ||
| def test_load_file_chunks(sample_dag, database_table_fixture): |
There was a problem hiding this comment.
Shouldn't we still support chunks in Postgres?
There was a problem hiding this comment.
@tatiana AFAICT it seems that since we're now using StringIO that the buffering is handled for us, I can investigate further though.
There was a problem hiding this comment.
Yes, please, @dimberman , I'd love to understand this.
|
Amazing results, @dimberman , some minor comments! One of the cases we're covering in the Snowflake optimization is if the user gives the Table columns - should we also cover this in the Postgres optimization? It could be in a separate PR. |
| with patch(type_method_map_fixture[FileType(filetype)]) as mocked_read, patch( | ||
| "astro.files.base.smart_open.open", mock_open(read_data=data) | ||
| ) as mocked_smart_open: |
There was a problem hiding this comment.
What is the benefit of declaring the patches within the test instead of using the decorator? It would be great if we were consistent with other tests.
There was a problem hiding this comment.
@tatiana from what I can tell mock_open must be used as part of the context statement. When I tried adding mock_open to the decorator it didn't work. I can try again.
There was a problem hiding this comment.
@tatiana is that ok? :)
There was a problem hiding this comment.
Hey @dimberman if the value of read_data didn't change, we could do something like this:
@patch( "astro.files.base.smart_open.open", new_callable=mock_open, read_data="data")
But since it changes, I believe it would add much complexity to have this fixture dependent on the filetype! So it's okay, thank you!
|
@tatiana can you point me to where this work is being done (re: allowing users to define their own schema)? |
Codecov Report
@@ Coverage Diff @@
## main #531 +/- ##
==========================================
+ Coverage 92.35% 92.43% +0.07%
==========================================
Files 40 40
Lines 1557 1573 +16
Branches 203 204 +1
==========================================
+ Hits 1438 1454 +16
- Misses 92 93 +1
+ Partials 27 26 -1
Continue to review full report at Codecov.
|
| @@ -46,7 +46,7 @@ def create_dag(database_name, table_args, dataset): | |||
| dag_name = f"load_file_{dataset_name}_into_{database_name}" | |||
|
|
|||
| with DAG(dag_name, schedule_interval=None, start_date=START_DATE) as dag: | |||
| chunk_size = int(os.getenv("ASTRO_CHUNK_SIZE", DEFAULT_CHUNK_SIZE)) | |||
| chunk_size = int(os.getenv("ASTRO_CHUNK_SIZE", str(DEFAULT_CHUNK_SIZE))) | |||
There was a problem hiding this comment.
@dimberman We should use airflow's config manager here.
| chunk_size = int(os.getenv("ASTRO_CHUNK_SIZE", str(DEFAULT_CHUNK_SIZE))) | |
| chunk_size = conf.get("astro_sdk", "chunk_size", str(DEFAULT_CHUNK_SIZE)) |
There was a problem hiding this comment.
@utkarsharma2 this isn't really an astro config. It's just an env variable we set for the benchmarking.
|
@dimberman Sorry for the delay!
|
Description
What is the current behavior?
Currently the postgres load numbers are far too slow for any production use-case. We need to find optimizations to bring the 10GB loading time to under 5 minutes for customers to reasonably use this feature.
Previous numbers:
closes: #428
What is the new behavior?
In this PR we found two primary optimizations that have led to significant speed up
The first optimization is to not directly place a smart_open into the
readfunction of pandas. There have been multiple reports of pandas and smart_open working very slow when used together. Instead, we read thesmart_openfile into an IO buffer (either StringIO or BytesIO). This also has the benefit that we do not need to write to disc if there is enough memory to hold the object.The second optimization is postgres specific. The
pandas.to_sqlfunction uses the HIGHLY suboptimalINSERTstatement for postgres. This is 10X slower than theCOPYcommand. By saving objects as CSV buffers, we can use COPY to load the data into postgres as a stream instead of one value at a time.Here are the new numbers with the optimizations:
Does this introduce a breaking change?
No
Checklist