Skip to content

Commit

Permalink
[#5453] [YSQL] Fix OOM on File Sourced YB Relations in COPY FROM Query
Browse files Browse the repository at this point in the history
Summary:
This change resets memory context regularly on all YB relations that are copied from files.
Instead of resetting memory context at batch size using `rows_per_transaction` query option, memory context will be reset per row when certain conditions hold. See below for details:

With this change:
1. Memory context resets per row instead of per batch size when memory is resettable.
  - eg. `COPY tab FROM 'absoluteFilePath'` will reset memory at every row when conditions are met.

2. Memory context reset per row occurs on YB relations (which excludes temporary tables) and when rows are read from file, not stdin.
  - if `cstate->filename != NULL`, file is consumed either directly or passed in using program option
  - if `cstate->filename == NULL`, query is consumed from stdin.

Note, large files read through stdin will run into OOM since old context path must be used to hold all rows in memory before inserting to table. That issue is tracked in [[ #5603 | here ]].

Test Plan:
Rebuilt and reran Java tests.
Manually verified performance on large queries in ysqlsh.

Reviewers: alex, mihnea, jason

Reviewed By: jason

Subscribers: yql

Differential Revision: https://phabricator.dev.yugabyte.com/D9313
  • Loading branch information
emhna committed Sep 5, 2020
1 parent ecf9a28 commit 7c7c841
Showing 1 changed file with 23 additions and 10 deletions.
33 changes: 23 additions & 10 deletions src/postgres/src/backend/commands/copy.c
Original file line number Diff line number Diff line change
Expand Up @@ -2391,7 +2391,8 @@ CopyFrom(CopyState cstate)
int prev_leaf_part_index = -1;
bool useNonTxnInsert;
bool isBatchTxnCopy = cstate->batch_size > 0;
MemoryContext batch_txn_copy_context;
bool shouldResetMemoryPerRow = IsYBRelation(cstate->rel) && cstate->filename != NULL;
MemoryContext row_context;

#define MAX_BUFFERED_TUPLES 1000
HeapTuple *bufferedTuples = NULL; /* initialize to silence warning */
Expand Down Expand Up @@ -2710,13 +2711,18 @@ CopyFrom(CopyState cstate)
errhint("Non-transactional COPY is not supported on relations with "
"secondary indices or triggers.")));

if (isBatchTxnCopy)
/*
* For YBRelations copied from stdin, require
* all rows to be held in memory before being inserted into relation.
* Thus, memory context will not be reset per row for stdin query types.
*/
if (shouldResetMemoryPerRow)
{
batch_txn_copy_context =
row_context =
AllocSetContextCreate(GetCurrentMemoryContext(),
"BATCH TXN COPY FROM",
"COPY FROM ROW CONTEXT",
ALLOCSET_DEFAULT_SIZES);
oldcontext = MemoryContextSwitchTo(batch_txn_copy_context);
oldcontext = MemoryContextSwitchTo(row_context);
}

bool has_more_tuples = true;
Expand Down Expand Up @@ -2747,7 +2753,7 @@ CopyFrom(CopyState cstate)
}

/* Switch into its memory context */
if (!isBatchTxnCopy)
if (!shouldResetMemoryPerRow)
MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));

has_more_tuples = NextCopyFrom(cstate, econtext, values, nulls, &loaded_oid);
Expand All @@ -2767,7 +2773,7 @@ CopyFrom(CopyState cstate)
tuple->t_tableOid = RelationGetRelid(resultRelInfo->ri_RelationDesc);

/* Triggers and stuff need to be invoked in query context. */
if (!isBatchTxnCopy)
if (!shouldResetMemoryPerRow)
MemoryContextSwitchTo(oldcontext);

/* Place tuple in tuple slot --- but slot shouldn't free it */
Expand Down Expand Up @@ -3009,16 +3015,23 @@ CopyFrom(CopyState cstate)
resultRelInfo = saved_resultRelInfo;
estate->es_result_relation_info = resultRelInfo;
}

/*
* Free context per row.
*/
if (shouldResetMemoryPerRow)
{
MemoryContextReset(row_context);
}
}
/*
* Commit transaction and free context per batch.
* Commit transaction per batch.
* When CopyFrom method is called, we are already inside a transaction block
* and relevant transaction state properties have been previously set.
*/
if (isBatchTxnCopy)
{
YBCCommitTransaction();
MemoryContextReset(batch_txn_copy_context);
YBInitializeTransaction();
}
}
Expand All @@ -3042,7 +3055,7 @@ CopyFrom(CopyState cstate)
*/
if (isBatchTxnCopy)
{
MemoryContextDelete(batch_txn_copy_context);
MemoryContextDelete(row_context);
}

/*
Expand Down

0 comments on commit 7c7c841

Please sign in to comment.