-
Notifications
You must be signed in to change notification settings - Fork 40
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Optimize postgres performance #531
Conversation
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why have we introduced this change?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
BTW: are there more bugs besides #515?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Shouldn't we still support chunks in Postgres?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, please, @dimberman , I'd love to understand this.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@tatiana is that ok? :)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@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
read
function of pandas. There have been multiple reports of pandas and smart_open working very slow when used together. Instead, we read thesmart_open
file 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_sql
function uses the HIGHLY suboptimalINSERT
statement for postgres. This is 10X slower than theCOPY
command. 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