Optimize SqlTableWriter

[kaapstorm]

Optimizes SqlTableWriter:

• Instead of running make_table_compatible() on every row, only runs on the first 10 (SCHEMA_CHECK_ROWS const) rows.
• Avoids increasing SQLAlchemy MetaData object growth.
• Upserts in batches of 1000 rows.
• Uses native bulk upserts when possible (PostgreSQL, MySQL) instead of one row at a time.
• Commits 1000 rows at a time instead of maintaining an open transaction for an entire export.

[millerdev]

Suggestion: convert to DRAFT until you're done pushing to the branch

[kaapstorm]

To make review easier, I'm going to split this into two PRs. I'll keep this one open, and chop it off after commit 7, Clean up SQL writer memory optimization spec.

[dannyroberts]

I had one substantive comment that I think will help a lot with the readability of one section, but otherwise looks good!

[dannyroberts]

If there are multiple rows in the same batch that have different schema mismatches, will the schema be updated for all rows in the batch that fail before retrying, or will it keep retrying until there are no more rows failing? The way it's written here it almost sounds like it'll try once, but if it fails on another row it'll just give up.

If the error we get back can only tell us about one row at a time, it almost seems better to just call make_table_compatible() for all 1000 rows and then retry once than to fix one, retry, and fail.

[dannyroberts]

Okay this answers my earlier question on the spec.

[dannyroberts]

Will this work if there are exactly SCHEMA_CHECK_ROWS rows? Looking at the rows after the loop, I think so? But it's just a tad hard to follow.

[dannyroberts]

I wonder if we can do something nice with itertools that makes this code flow more naturally. Right now there's a lot of iteration accounting with each step doing conditional cleanup for the last step and/or lookahead for the next in a way that's a bit hard to follow.

I'm thinking something like this:

if table_spec.rows:
    # ... create_table(...) ...

row_stream = iter(table_spec.rows)

for row in itertools.islice(row_stream, SCHEMA_CHECK_ROWS):
    table = self.make_table_compatible...
    self.upsert...

logger.debug("Schema check complete...

for batch in itertools.batched(row_stream, BATCH_SIZE):
    self._flush_batch(table, batch, ...

I think that's it—I think there'd then be no cleanup/commit/flush necessary afterwards.

[kaapstorm]

So much better! f7fba6e

[jingcheng16]

No blocking comments.

[jingcheng16]

Is this all possible exceptions when there is a incompatibility?

[kaapstorm]

d3f6566

[jingcheng16]

Can we be explicit that this is just a schema check for the first SCHEMA_CHECK_ROWS rows? This might be helpful for people who didn't read the code and just read the log.

[kaapstorm]

dd54c0f

[jingcheng16]

I find this set comprehension hard to read...maybe it is just me...

[kaapstorm]

804873e

[jingcheng16]

AI comment:
The old upsert() (writers.py:585) does:

   row_dict = {col: val for col, val in row_dict.items() if val is not None}                                                                                                                                  
   # then UPDATE ... SET (only the non-None cols)    

So an existing row's columns that aren't in row_dict are preserved.

The new bulk_upsert() builds batch_keys as the union of non-None keys across the whole batch, then in PG iterates stmt.excluded and in MySQL iterates table.columns to construct the SET clause. Worth verifying whether columns that exist in the table but were dropped from batch_keys (because every row in that batch had None for them) end up getting clobbered to NULL on update. Old code wouldn't have touched them; new code might.

Thoughts? Sounds like a real concern

[kaapstorm]

Nicely spotted! ca0221e

[dannyroberts]

Very nice! The itertools refactor is so very much easier to follow! Left one small thought on the same topic, but totally up to you—looks good as is.

[dannyroberts]

Nit: Would the following work? I think the only difference would be that self.make_table_compatible would not be called on the first row, but that might be okay (just based on the names of the functions) if create_table is already taking into account first_row's schema? I guess technically we'd also need SCHEMA_CHECK_ROWS - 1 below. What I like about this is that then we're literally doing one full clean sweep through the data, no lookahead or backtracking, so there's no data flow to keep in your head. Process first row. Process next N rows. Process rest.

Not a big deal at all; the way you have it is not much harder to understand.

Suggested change

	rows = (dict(zip(headings, row)) for row in table_spec.rows)
	first_row = next(rows, None)
	if first_row is None:
	return
	row_stream = itertools.chain([first_row], rows)
	table = self.get_table(table_name)
	if table is None:
	table = self.create_table(table_name, first_row, data_type_dict)
	row_stream = (dict(zip(headings, row)) for row in table_spec.rows)
	first_row = next(row_stream, None)
	if first_row is None:
	return
	table = self.get_table(table_name)
	if table is None:
	table = self.create_table(table_name, first_row, data_type_dict)
	self.upsert(table, first_row)

[kaapstorm]

I think the only difference would be that self.make_table_compatible would not be called on the first row.

Yes, and neither would self.upsert (on the next line), so the first row would not end up in the database.