JSON_TABLE: Add support for NESTED paths and columns

A NESTED path allows to extract data from nested levels of JSON
objects given by the parent path expression, which are projected as
columns specified using a nested COLUMNS clause, just like the parent
COLUMNS clause.  Rows comprised from a NESTED columns are "joined"
to the row comprised from the parent columns.  If a particular NESTED
path evaluates to 0 rows, then the nested COLUMNS will emit NULLs,
making it an OUTER join.

NESTED columns themselves may include NESTED paths to allow
extracting data from arbitrary nesting levels, which are likewise
joined against the rows at the parent level.

Multiple NESTED paths at a given level are called "sibling" paths
and their rows are combined by UNIONing them, that is, after being
joined against the parent row as described above.

Author: Nikita Glukhov <n.gluhov@postgrespro.ru>
Author: Teodor Sigaev <teodor@sigaev.ru>
Author: Oleg Bartunov <obartunov@gmail.com>
Author: Alexander Korotkov <aekorotkov@gmail.com>
Author: Andrew Dunstan <andrew@dunslane.net>
Author: Amit Langote <amitlangote09@gmail.com>
Author: Jian He <jian.universality@gmail.com>

Reviewers have included (in no particular order):

Andres Freund, Alexander Korotkov, Pavel Stehule, Andrew Alsup,
Erik Rijkers, Zihong Yu, Himanshu Upadhyaya, Daniel Gustafsson,
Justin Pryzby, Álvaro Herrera, Jian He

Discussion: https://postgr.es/m/cd0bb935-0158-78a7-08b5-904886deac4b@postgrespro.ru
Discussion: https://postgr.es/m/20220616233130.rparivafipt6doj3@alap3.anarazel.de
Discussion: https://postgr.es/m/abd9b83b-aa66-f230-3d6d-734817f0995d%40postgresql.org
Discussion: https://postgr.es/m/CA+HiwqE4XTdfb1nW=Ojoy_tQSRhYt-q_kb6i5d4xcKyrLC1Nbg@mail.gmail.com

1 files changed, 162 insertions, 6 deletions

diff --git a/src/backend/utils/adt/jsonpath_exec.c b/src/backend/utils/adt/jsonpath_exec.c
index 75c468bc085..103572ed932 100644
--- a/src/backend/utils/adt/jsonpath_exec.c
+++ b/src/backend/utils/adt/jsonpath_exec.c
@@ -202,6 +202,18 @@ typedef struct JsonTablePlanState
 
 	/* Counter for ORDINAL columns */
 	int			ordinal;
+
+	/* Nested plan, if any */
+	struct JsonTablePlanState *nested;
+
+	/* Left sibling, if any */
+	struct JsonTablePlanState *left;
+
+	/* Right sibling, if any */
+	struct JsonTablePlanState *right;
+
+	/* Parent plan, if this is a nested plan */
+	struct JsonTablePlanState *parent;
 } JsonTablePlanState;
 
 /* Random number to identify JsonTableExecContext for sanity checking */
@@ -213,6 +225,12 @@ typedef struct JsonTableExecContext
 
 	/* State of the plan providing a row evaluated from "root" jsonpath */
 	JsonTablePlanState *rootplanstate;
+
+	/*
+	 * Per-column JsonTablePlanStates for all columns including the nested
+	 * ones.
+	 */
+	JsonTablePlanState **colplanstates;
 } JsonTableExecContext;
 
 /* strict/lax flags is decomposed into four [un]wrap/error flags */
@@ -337,6 +355,7 @@ static void checkTimezoneIsUsedForCast(bool useTz, const char *type1,
 static void JsonTableInitOpaque(TableFuncScanState *state, int natts);
 static JsonTablePlanState *JsonTableInitPlan(JsonTableExecContext *cxt,
 											 JsonTablePlan *plan,
+											 JsonTablePlanState *parentstate,
 											 List *args,
 											 MemoryContext mcxt);
 static void JsonTableSetDocument(TableFuncScanState *state, Datum value);
@@ -345,6 +364,9 @@ static bool JsonTableFetchRow(TableFuncScanState *state);
 static Datum JsonTableGetValue(TableFuncScanState *state, int colnum,
 							   Oid typid, int32 typmod, bool *isnull);
 static void JsonTableDestroyOpaque(TableFuncScanState *state);
+static bool JsonTablePlanScanNextRow(JsonTablePlanState *planstate);
+static void JsonTableResetNestedPlan(JsonTablePlanState *planstate);
+static bool JsonTablePlanJoinNextRow(JsonTablePlanState *planstate);
 static bool JsonTablePlanNextRow(JsonTablePlanState *planstate);
 
 const TableFuncRoutine JsonbTableRoutine =
@@ -4087,8 +4109,14 @@ JsonTableInitOpaque(TableFuncScanState *state, int natts)
 		}
 	}
 
-	/* Initialize plan */
-	cxt->rootplanstate = JsonTableInitPlan(cxt, rootplan, args,
+	cxt->colplanstates = palloc(sizeof(JsonTablePlanState *) *
+								list_length(tf->colvalexprs));
+
+	/*
+	 * Initialize plan for the root path and, recursively, also any child
+	 * plans that compute the NESTED paths.
+	 */
+	cxt->rootplanstate = JsonTableInitPlan(cxt, rootplan, NULL, args,
 										   CurrentMemoryContext);
 
 	state->opaque = cxt;
@@ -4113,19 +4141,22 @@ JsonTableDestroyOpaque(TableFuncScanState *state)
 /*
  * JsonTableInitPlan
  *		Initialize information for evaluating jsonpath in the given
- *		JsonTablePlan
+ *		JsonTablePlan and, recursively, in any child plans
  */
 static JsonTablePlanState *
 JsonTableInitPlan(JsonTableExecContext *cxt, JsonTablePlan *plan,
+				  JsonTablePlanState *parentstate,
 				  List *args, MemoryContext mcxt)
 {
 	JsonTablePlanState *planstate = palloc0(sizeof(*planstate));
 
 	planstate->plan = plan;
+	planstate->parent = parentstate;
 
 	if (IsA(plan, JsonTablePathScan))
 	{
 		JsonTablePathScan *scan = (JsonTablePathScan *) plan;
+		int			i;
 
 		planstate->path = DatumGetJsonPathP(scan->path->value->constvalue);
 		planstate->args = args;
@@ -4135,6 +4166,21 @@ JsonTableInitPlan(JsonTableExecContext *cxt, JsonTablePlan *plan,
 		/* No row pattern evaluated yet. */
 		planstate->current.value = PointerGetDatum(NULL);
 		planstate->current.isnull = true;
+
+		for (i = scan->colMin; i >= 0 && i <= scan->colMax; i++)
+			cxt->colplanstates[i] = planstate;
+
+		planstate->nested = scan->child ?
+			JsonTableInitPlan(cxt, scan->child, planstate, args, mcxt) : NULL;
+	}
+	else if (IsA(plan, JsonTableSiblingJoin))
+	{
+		JsonTableSiblingJoin *join = (JsonTableSiblingJoin *) plan;
+
+		planstate->left = JsonTableInitPlan(cxt, join->lplan, parentstate,
+											args, mcxt);
+		planstate->right = JsonTableInitPlan(cxt, join->rplan, parentstate,
+											 args, mcxt);
 	}
 
 	return planstate;
@@ -4193,16 +4239,56 @@ JsonTableResetRowPattern(JsonTablePlanState *planstate, Datum item)
 }
 
 /*
- * Fetch next row from a JsonTablePlan's path evaluation result.
+ * Fetch next row from a JsonTablePlan.
  *
  * Returns false if the plan has run out of rows, true otherwise.
  */
 static bool
 JsonTablePlanNextRow(JsonTablePlanState *planstate)
 {
-	JsonbValue *jbv = JsonValueListNext(&planstate->found, &planstate->iter);
+	if (IsA(planstate->plan, JsonTablePathScan))
+		return JsonTablePlanScanNextRow(planstate);
+	else if (IsA(planstate->plan, JsonTableSiblingJoin))
+		return JsonTablePlanJoinNextRow(planstate);
+	else
+		elog(ERROR, "invalid JsonTablePlan %d", (int) planstate->plan->type);
+
+	Assert(false);
+	/* Appease compiler */
+	return false;
+}
+
+/*
+ * Fetch next row from a JsonTablePlan's path evaluation result and from
+ * any child nested path(s).
+ *
+ * Returns true if any of the paths (this or the nested) has more rows to
+ * return.
+ *
+ * By fetching the nested path(s)'s rows based on the parent row at each
+ * level, this essentially joins the rows of different levels.  If a nested
+ * path at a given level has no matching rows, the columns of that level will
+ * compute to NULL, making it an OUTER join.
+ */
+static bool
+JsonTablePlanScanNextRow(JsonTablePlanState *planstate)
+{
+	JsonbValue *jbv;
 	MemoryContext oldcxt;
 
+	/*
+	 * If planstate already has an active row and there is a nested plan,
+	 * check if it has an active row to join with the former.
+	 */
+	if (!planstate->current.isnull)
+	{
+		if (planstate->nested && JsonTablePlanNextRow(planstate->nested))
+			return true;
+	}
+
+	/* Fetch new row from the list of found values to set as active. */
+	jbv = JsonValueListNext(&planstate->found, &planstate->iter);
+
 	/* End of list? */
 	if (jbv == NULL)
 	{
@@ -4223,6 +4309,76 @@ JsonTablePlanNextRow(JsonTablePlanState *planstate)
 	/* Next row! */
 	planstate->ordinal++;
 
+	/* Process nested plan(s), if any. */
+	if (planstate->nested)
+	{
+		/* Re-evaluate the nested path using the above parent row. */
+		JsonTableResetNestedPlan(planstate->nested);
+
+		/*
+		 * Now fetch the nested plan's current row to be joined against the
+		 * parent row.  Any further nested plans' paths will be re-evaluated
+		 * reursively, level at a time, after setting each nested plan's
+		 * current row.
+		 */
+		(void) JsonTablePlanNextRow(planstate->nested);
+	}
+
+	/* There are more rows. */
+	return true;
+}
+
+/*
+ * Re-evaluate the row pattern of a nested plan using the new parent row
+ * pattern.
+ */
+static void
+JsonTableResetNestedPlan(JsonTablePlanState *planstate)
+{
+	/* This better be a child plan. */
+	Assert(planstate->parent != NULL);
+	if (IsA(planstate->plan, JsonTablePathScan))
+	{
+		JsonTablePlanState *parent = planstate->parent;
+
+		if (!parent->current.isnull)
+			JsonTableResetRowPattern(planstate, parent->current.value);
+
+		/*
+		 * If this plan itself has a child nested plan, it will be reset when
+		 * the caller calls JsonTablePlanNextRow() on this plan.
+		 */
+	}
+	else if (IsA(planstate->plan, JsonTableSiblingJoin))
+	{
+		JsonTableResetNestedPlan(planstate->left);
+		JsonTableResetNestedPlan(planstate->right);
+	}
+}
+
+/*
+ * Fetch the next row from a JsonTableSiblingJoin.
+ *
+ * This is essentially a UNION between the rows from left and right siblings.
+ */
+static bool
+JsonTablePlanJoinNextRow(JsonTablePlanState *planstate)
+{
+
+	/* Fetch row from left sibling. */
+	if (!JsonTablePlanNextRow(planstate->left))
+	{
+		/*
+		 * Left sibling ran out of rows, so start fetching from the right
+		 * sibling.
+		 */
+		if (!JsonTablePlanNextRow(planstate->right))
+		{
+			/* Right sibling ran out of row, so there are more rows. */
+			return false;
+		}
+	}
+
 	return true;
 }
 
@@ -4256,7 +4412,7 @@ JsonTableGetValue(TableFuncScanState *state, int colnum,
 		GetJsonTableExecContext(state, "JsonTableGetValue");
 	ExprContext *econtext = state->ss.ps.ps_ExprContext;
 	ExprState  *estate = list_nth(state->colvalexprs, colnum);
-	JsonTablePlanState *planstate = cxt->rootplanstate;
+	JsonTablePlanState *planstate = cxt->colplanstates[colnum];
 	JsonTablePlanRowSource *current = &planstate->current;
 	Datum		result;