fix: Cannot use entity functions in list comprehension

src/backend/executor/execExpr.c

@@ -4404,8 +4404,24 @@ ExecInitCypherListComp(ExprEvalStep *scratch, CypherListCompExpr *listcompexpr,
 	null_stepno = state->steps_len - 1;
 
 	begin_step.opcode = EEOP_CYPHERLISTCOMP_BEGIN;
-	begin_step.d.cypherlistcomp.liststate =
-		(JsonbParseState **) palloc(sizeof(JsonbParseState *));
+	begin_step.d.cypherlistcomp.result_type =
+		(Oid *) palloc(sizeof(Oid));
+	*begin_step.d.cypherlistcomp.result_type =
+		exprType((Node *) listcompexpr);
+
+	/* Only allocate the state pointer we'll actually use */
+	if (*begin_step.d.cypherlistcomp.result_type == JSONBOID)
+	{
+		begin_step.d.cypherlistcomp.jstate =
+			(JsonbParseState **) palloc(sizeof(JsonbParseState *));
+		begin_step.d.cypherlistcomp.astate = NULL;
+	}
+	else
+	{
+		begin_step.d.cypherlistcomp.astate =
+			(ArrayBuildState **) palloc(sizeof(ArrayBuildState *));
+		begin_step.d.cypherlistcomp.jstate = NULL;
+	}
 	ExprEvalPushStep(state, &begin_step);
 
 	init_step.opcode = EEOP_CYPHERLISTCOMP_ITER_INIT;
@@ -4480,17 +4496,25 @@ ExecInitCypherListComp(ExprEvalStep *scratch, CypherListCompExpr *listcompexpr,
 	elem_step.opcode = EEOP_CYPHERLISTCOMP_ELEM;
 	elem_step.d.cypherlistcomp.elemvalue = elem_resvalue;
 	elem_step.d.cypherlistcomp.elemnull = elem_resnull;
-	elem_step.d.cypherlistcomp.liststate =
-		begin_step.d.cypherlistcomp.liststate;
+	elem_step.d.cypherlistcomp.astate =
+		begin_step.d.cypherlistcomp.astate;
+	elem_step.d.cypherlistcomp.jstate =
+		begin_step.d.cypherlistcomp.jstate;
+	elem_step.d.cypherlistcomp.result_type =
+		begin_step.d.cypherlistcomp.result_type;
 	ExprEvalPushStep(state, &elem_step);
 
 	loop_step.opcode = EEOP_JUMP;
 	loop_step.d.jump.jumpdone = next_stepno;
 	ExprEvalPushStep(state, &loop_step);
 
 	scratch->opcode = EEOP_CYPHERLISTCOMP_END;
-	scratch->d.cypherlistcomp.liststate =
-		begin_step.d.cypherlistcomp.liststate;
+	scratch->d.cypherlistcomp.astate =
+		begin_step.d.cypherlistcomp.astate;
+	scratch->d.cypherlistcomp.jstate =
+		begin_step.d.cypherlistcomp.jstate;
+	scratch->d.cypherlistcomp.result_type =
+		begin_step.d.cypherlistcomp.result_type;
 	ExprEvalPushStep(state, scratch);
 	end_stepno = state->steps_len - 1;
 

src/backend/executor/execExprInterp.c

@@ -5534,54 +5534,122 @@ cypher_access_index(CypherIndexResult *cidxres, const int nelems)
 	return -1;
 }
 
+/*
+ * ExecEvalCypherListCompBegin
+ *		Initialize array building state for list comprehension.
+ *
+ * For JSONB results, initializes JsonbParseState with array container.
+ * For typed arrays (vertex/edge), uses ArrayBuildState.
+ */
 void
 ExecEvalCypherListCompBegin(ExprState *state, ExprEvalStep *op)
 {
-	*op->d.cypherlistcomp.liststate = NULL;
-	pushJsonbValue(op->d.cypherlistcomp.liststate,
-				   WJB_BEGIN_ARRAY, NULL);
+	Oid			result_type = *op->d.cypherlistcomp.result_type;
+
+	if (result_type == JSONBOID)
+	{
+		*op->d.cypherlistcomp.jstate = NULL;
+		pushJsonbValue(op->d.cypherlistcomp.jstate,
+					   WJB_BEGIN_ARRAY, NULL);
+	}
+	else
+	{
+		Oid			elem_type;
+
+		elem_type = get_element_type(result_type);
+		if (elem_type == InvalidOid)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("cannot get element type for list comprehension")));
+
+		*op->d.cypherlistcomp.astate = initArrayResult(elem_type,
+													   CurrentMemoryContext,
+													   false);
+	}
 }
 
+/*
+ * ExecEvalCypherListCompElem
+ *		Accumulate one element into the list comprehension result.
+ *
+ * For jsonb, converts element to JsonbValue and pushes to JsonbParseState.
+ * For typed arrays, uses accumArrayResult to append to ArrayBuildState.
+ */
 void
 ExecEvalCypherListCompElem(ExprState *state, ExprEvalStep *op)
 {
-	JsonbValue	_ejv;
-	JsonbValue *ejv;
+	Datum		elemvalue = *op->d.cypherlistcomp.elemvalue;
+	bool		elemnull = *op->d.cypherlistcomp.elemnull;
+	Oid			result_type = *op->d.cypherlistcomp.result_type;
 
-	if (*op->d.cypherlistcomp.elemnull)
+	if (result_type == JSONBOID)
 	{
-		_ejv.type = jbvNull;
-		ejv = &_ejv;
-	}
-	else
-	{
-		Jsonb	   *ejb;
+		JsonbValue	_ejv;
+		JsonbValue *ejv;
 
-		ejb = DatumGetJsonbP(*op->d.cypherlistcomp.elemvalue);
-		if (JB_ROOT_IS_SCALAR(ejb))
+		if (elemnull)
 		{
-			ejv = getIthJsonbValueFromContainer(&ejb->root, 0);
+			_ejv.type = jbvNull;
+			ejv = &_ejv;
 		}
 		else
 		{
-			_ejv.type = jbvBinary;
-			_ejv.val.binary.data = &ejb->root;
-			ejv = &_ejv;
+			Jsonb	   *ejb = DatumGetJsonbP(elemvalue);
+
+			if (JB_ROOT_IS_SCALAR(ejb))
+			{
+				ejv = getIthJsonbValueFromContainer(&ejb->root, 0);
+			}
+			else
+			{
+				_ejv.type = jbvBinary;
+				_ejv.val.binary.data = &ejb->root;
+				ejv = &_ejv;
+			}
 		}
-	}
 
-	pushJsonbValue(op->d.cypherlistcomp.liststate, WJB_ELEM, ejv);
+		pushJsonbValue(op->d.cypherlistcomp.jstate, WJB_ELEM, ejv);
+	}
+	else
+	{
+		ArrayBuildState *astate = *op->d.cypherlistcomp.astate;
+		Oid			elem_type = astate->element_type;
+
+		*op->d.cypherlistcomp.astate = accumArrayResult(astate,
+														elemvalue,
+														elemnull,
+														elem_type,
+														CurrentMemoryContext);
+	}
 }
 
+/*
+ * ExecEvalCypherListCompEnd
+ *		Finalize the list comprehension result into a complete array.
+ *
+ * For jsonb, finalizes JsonbParseState into a complete JSONB array.
+ * For typed arrays, calls makeArrayResult to build the final typed array.
+ */
 void
 ExecEvalCypherListCompEnd(ExprState *state, ExprEvalStep *op)
 {
-	JsonbValue *jv;
+	Oid			result_type = *op->d.cypherlistcomp.result_type;
+
+	if (result_type == JSONBOID)
+	{
+		JsonbValue *jv;
 
-	jv = pushJsonbValue(op->d.cypherlistcomp.liststate,
-						WJB_END_ARRAY, NULL);
+		jv = pushJsonbValue(op->d.cypherlistcomp.jstate,
+							WJB_END_ARRAY, NULL);
+		*op->resvalue = JsonbPGetDatum(JsonbValueToJsonb(jv));
+	}
+	else
+	{
+		ArrayBuildState *astate = *op->d.cypherlistcomp.astate;
+
+		*op->resvalue = makeArrayResult(astate, CurrentMemoryContext);
+	}
 
-	*op->resvalue = JsonbPGetDatum(JsonbValueToJsonb(jv));
 	*op->resnull = false;
 }
 
@@ -5653,9 +5721,11 @@ ExecEvalCypherListCompIterInitNext(ExprState *state, ExprEvalStep *op)
 	}
 	else
 	{
-		CypherListCompArrayIterator *array_iter = op->d.cypherlistcomp_iter.array_iterator;
+		CypherListCompArrayIterator *array_iter;
 		Datum		vertex_or_edge_datum;
 
+		array_iter = op->d.cypherlistcomp_iter.array_iterator;
+
 		if (array_iter->array_position >= array_iter->array_size)
 		{
 			*op->resvalue = (Datum) 0;
@@ -5674,15 +5744,7 @@ ExecEvalCypherListCompIterInitNext(ExprState *state, ExprEvalStep *op)
 		{
 			return;
 		}
-
-		if (array_iter->array_typid == VERTEXOID)
-		{
-			*op->resvalue = getVertexPropDatum(vertex_or_edge_datum);
-		}
-		else
-		{
-			*op->resvalue = getEdgePropDatum(vertex_or_edge_datum);
-		}
+		*op->resvalue = vertex_or_edge_datum;
 	}
 }
 

src/backend/nodes/nodeFuncs.c

@@ -280,10 +280,49 @@ exprType(const Node *expr)
 			type = JSONBOID;
 			break;
 		case T_CypherListCompExpr:
-			type = JSONBOID;
-			break;
+		{
+			/*
+			 * Determine result type of list comprehension:
+			 * - Identity case, same as input list type
+			 * - Graph object elem, appropriate typed array (vertex/edge)
+			 * - Other elem, jsonb array
+			 */
+			CypherListCompExpr *clcexpr = (CypherListCompExpr *) expr;
+			bool		is_identity;
+			Oid			elem_type;
+
+			is_identity = (clcexpr->elem == NULL) ||
+						  (clcexpr->cond == NULL &&
+						   IsA(clcexpr->elem, CypherListCompVar));
+
+			if (is_identity)
+			{
+				type = exprType((Node *) clcexpr->list);
+			}
+			else if (clcexpr->elem != NULL)
+			{
+				/* Check the elem's result type */
+				elem_type = exprType((Node *) clcexpr->elem);
+
+				if (elem_type == VERTEXOID)
+					type = VERTEXARRAYOID;
+				else if (elem_type == EDGEOID)
+					type = EDGEARRAYOID;
+				else
+					type = JSONBOID;
+			}
+			else
+			{
+				/* No elem, return jsonb array */
+				type = JSONBOID;
+			}
+		}
+		break;
 		case T_CypherListCompVar:
-			type = JSONBOID;
+			{
+				CypherListCompVar *clcvar = (CypherListCompVar *) expr;
+				type = clcvar->elem_type;
+			}
 			break;
 		case T_CypherAccessExpr:
 			type = JSONBOID;

src/backend/parser/parse_cypher_expr.c

@@ -444,6 +444,7 @@ transformListCompColumnRef(ParseState *pstate, ColumnRef *cref, char *varname)
 
 	clcvar = makeNode(CypherListCompVar);
 	clcvar->varname = pstrdup(varname);
+	clcvar->elem_type = pstate->p_lc_elem_type;
 	clcvar->location = cref->location;
 
 	if (list_length(cref->fields) > 1)
@@ -723,30 +724,83 @@ transformCypherListComp(ParseState *pstate, CypherListComp *clc)
 	Node	   *list;
 	Oid			type;
 	char	   *save_varname;
+	Oid			save_elem_type;
 	Node	   *cond;
 	Node	   *elem;
 	CypherListCompExpr *clcexpr;
 
 	list = transformCypherExprRecurse(pstate, (Node *) clc->list);
 	type = exprType(list);
 
+	/*
+	 * Determine the element type from the list's array type.
+	 * This is used for type checking of the iteration variable in LC.
+	 */
 	switch (type)
 	{
 		case JSONBOID:
+			pstate->p_lc_elem_type = JSONBOID;
+			break;
 		case VERTEXARRAYOID:
+			pstate->p_lc_elem_type = VERTEXOID;
+			break;
 		case EDGEARRAYOID:
+			pstate->p_lc_elem_type = EDGEOID;
 			break;
 		default:
 			list = coerce_all_to_jsonb(pstate, list);
+			pstate->p_lc_elem_type = JSONBOID;
 			break;
 	}
 
 	save_varname = pstate->p_lc_varname;
+	save_elem_type = pstate->p_lc_elem_type;
 	pstate->p_lc_varname = clc->varname;
 	cond = transformCypherWhere(pstate, clc->cond, pstate->p_expr_kind);
 	elem = transformCypherExprRecurse(pstate, clc->elem);
+
+	/*
+	 * If elem is NULL, implicitly use the iteration variable.
+	 */
+	if (elem == NULL && clc->varname != NULL)
+	{
+		CypherListCompVar *clcvar = makeNode(CypherListCompVar);
+		clcvar->varname = pstrdup(clc->varname);
+		clcvar->elem_type = pstate->p_lc_elem_type;
+		clcvar->location = clc->location;
+		elem = (Node *) clcvar;
+	}
+
 	pstate->p_lc_varname = save_varname;
-	elem = coerce_to_jsonb(pstate, elem, "list comprehension result");
+	pstate->p_lc_elem_type = save_elem_type;
+
+	/*
+	 * Determine if coercion to JSONB is needed for the result expression.
+	 * If it is identity case (the iteration variable itself) or a graph
+	 * object (vertex/edge), no coercion is needed.
+	 * Otherwise, coerce it to JSONB.
+	 */
+	if (elem != NULL)
+	{
+		bool		is_identity_elem;
+		bool		is_graph_object_result;
+		Oid			elem_result_type;
+
+		/* Check if this is the identity case */
+		is_identity_elem = (cond == NULL) &&
+						   (IsA(elem, CypherListCompVar));
+
+		/* Check if elem returns a graph object (vertex/edge) */
+		elem_result_type = exprType(elem);
+		is_graph_object_result = (elem_result_type == VERTEXOID ||
+								  elem_result_type == EDGEOID);
+
+		if (!is_identity_elem && !is_graph_object_result)
+		{
+			elem = coerce_to_jsonb(pstate, elem,
+								   "list comprehension result");
+		}
+	}
 
 	clcexpr = makeNode(CypherListCompExpr);
 	clcexpr->list = (Expr *) list;

src/include/executor/execExpr.h

@@ -16,6 +16,7 @@
 
 #include "executor/nodeAgg.h"
 #include "nodes/execnodes.h"
+#include "utils/array.h"
 #include "utils/jsonb.h"
 #include "utils/arrayaccess.h"
 
@@ -731,7 +732,9 @@ typedef struct ExprEvalStep
 		{
 			Datum	   *elemvalue;
 			bool	   *elemnull;
-			JsonbParseState **liststate;
+			ArrayBuildState **astate; /* array build state */
+			JsonbParseState **jstate; /* jsonb parse state */
+			Oid		   *result_type; /* result array type */
 		}			cypherlistcomp;
 
 		struct