v5-0004-jsonb-optimize-multi-subscript-casts-via-extract-.patch

From 1c5bc428388e6a4f97c1e684b7705b98e5d0baf9 Mon Sep 17 00:00:00 2001
From: Haibo Yan <haibo.yan@apple.com>
Date: Sun, 26 Apr 2026 18:49:33 -0700
Subject: [PATCH v5 4/5] jsonb: optimize multi-subscript casts via extract-path
 rewrite

Extend the existing support-function rewrite to multi-subscript jsonb
subscripting chains, lowering them to the extract-path typed extractor
family already introduced by the series.

Only chains whose integer subscripts are all constants are eligible for
the rewrite; non-constant integer subscripts (which would require a
runtime CoerceViaIO conversion) cause the entire chain to be left in
its original form.

This keeps single-subscript behavior unchanged and supports casts to
numeric, bool, int2, int4, int8, float4, and float8.
---
 src/backend/utils/adt/jsonb.c       | 116 +++++++++++++++++---
 src/test/regress/expected/jsonb.out | 157 ++++++++++++++++++++++++++++
 src/test/regress/sql/jsonb.sql      |  43 ++++++++
 3 files changed, 301 insertions(+), 15 deletions(-)

diff --git a/src/backend/utils/adt/jsonb.c b/src/backend/utils/adt/jsonb.c
index 7eba437b8a4..d5c0871a213 100644
--- a/src/backend/utils/adt/jsonb.c
+++ b/src/backend/utils/adt/jsonb.c
@@ -1835,6 +1835,7 @@ cannotCastJsonbValue(enum jbvType type, const char *sqltype, Node *escontext)
  *   - jsonb_object_field(j, 'key')  /  j -> 'key'  /  j['key']
  *   - jsonb_array_element(j, idx)   /  j -> idx    /  j[idx]
  *   - jsonb_extract_path(j, ...)    /  j #> '{a,b}'
+ *   - multi-subscript chains         j['a']['b'], j['a'][0], etc.
  */
 Datum
 jsonb_cast_support(PG_FUNCTION_ARGS)
@@ -1885,35 +1886,120 @@ jsonb_cast_support(PG_FUNCTION_ARGS)
 		else if (IsA(arg, SubscriptingRef))
 		{
 			SubscriptingRef *sbsref = (SubscriptingRef *) arg;
-			Node	   *subscript;
-			Oid			subscript_type;
+			int			nsubscripts;
 
 			/*
-			 * Handle single-subscript jsonb access with no slice and no
-			 * assignment.  Text subscripts map to object-field extraction;
-			 * int4 subscripts map to array-element extraction.
+			 * Handle jsonb subscript access with no slice and no assignment.
+			 * Single subscripts map to object-field or array-element
+			 * extraction; multi-subscript chains lower to the extract-path
+			 * family.
 			 */
 			if (sbsref->refcontainertype != JSONBOID)
 				PG_RETURN_POINTER(NULL);
-			if (list_length(sbsref->refupperindexpr) != 1)
-				PG_RETURN_POINTER(NULL);
 			if (sbsref->reflowerindexpr != NIL)
 				PG_RETURN_POINTER(NULL);
 			if (sbsref->refassgnexpr != NULL)
 				PG_RETURN_POINTER(NULL);
 
-			subscript = (Node *) linitial(sbsref->refupperindexpr);
-			subscript_type = exprType(subscript);
+			nsubscripts = list_length(sbsref->refupperindexpr);
 
-			if (subscript_type == TEXTOID)
+			if (nsubscripts == 1)
 			{
-				inner_funcid = F_JSONB_OBJECT_FIELD;
-				inner_args = list_make2(sbsref->refexpr, subscript);
+				/*
+				 * Single subscript: text maps to object-field, int4 maps to
+				 * array-element.
+				 */
+				Node	   *subscript;
+				Oid			subscript_type;
+
+				subscript = (Node *) linitial(sbsref->refupperindexpr);
+				subscript_type = exprType(subscript);
+
+				if (subscript_type == TEXTOID)
+				{
+					inner_funcid = F_JSONB_OBJECT_FIELD;
+					inner_args = list_make2(sbsref->refexpr, subscript);
+				}
+				else if (subscript_type == INT4OID)
+				{
+					inner_funcid = F_JSONB_ARRAY_ELEMENT;
+					inner_args = list_make2(sbsref->refexpr, subscript);
+				}
+				else
+					PG_RETURN_POINTER(NULL);
 			}
-			else if (subscript_type == INT4OID)
+			else if (nsubscripts >= 2)
 			{
-				inner_funcid = F_JSONB_ARRAY_ELEMENT;
-				inner_args = list_make2(sbsref->refexpr, subscript);
+				/*
+				 * Multi-subscript chain: build a text[] path and lower to
+				 * the extract-path family.  Each subscript must be text or
+				 * a constant int4; non-constant integer subscripts cause
+				 * the entire chain to be left unoptimized.
+				 */
+				List	   *path_elems = NIL;
+				ListCell   *lc;
+				ArrayExpr  *aexpr;
+
+				foreach(lc, sbsref->refupperindexpr)
+				{
+					Node	   *subscript = (Node *) lfirst(lc);
+					Oid			subscript_type = exprType(subscript);
+
+					if (subscript_type == TEXTOID)
+					{
+						path_elems = lappend(path_elems, subscript);
+					}
+					else if (subscript_type == INT4OID)
+					{
+						Const	   *con;
+
+						/*
+						 * Only constant integer subscripts can be safely
+						 * converted to text at plan time.  Non-constant
+						 * ones would require a runtime CoerceViaIO node;
+						 * decline the rewrite for the entire chain.
+						 */
+						if (!IsA(subscript, Const))
+							PG_RETURN_POINTER(NULL);
+
+						con = (Const *) subscript;
+
+						if (con->constisnull)
+						{
+							path_elems = lappend(path_elems,
+												 makeNullConst(TEXTOID,
+															   -1,
+															   InvalidOid));
+						}
+						else
+						{
+							char	   *str;
+
+							str = DatumGetCString(
+								DirectFunctionCall1(int4out,
+													con->constvalue));
+							path_elems = lappend(path_elems,
+								makeConst(TEXTOID, -1, InvalidOid, -1,
+										  CStringGetTextDatum(str),
+										  false, false));
+						}
+					}
+					else
+						PG_RETURN_POINTER(NULL);
+				}
+
+				aexpr = makeNode(ArrayExpr);
+				aexpr->array_typeid = TEXTARRAYOID;
+				aexpr->array_collid = InvalidOid;
+				aexpr->element_typeid = TEXTOID;
+				aexpr->elements = path_elems;
+				aexpr->multidims = false;
+				aexpr->list_start = -1;
+				aexpr->list_end = -1;
+				aexpr->location = -1;
+
+				inner_funcid = F_JSONB_EXTRACT_PATH;
+				inner_args = list_make2(sbsref->refexpr, aexpr);
 			}
 			else
 				PG_RETURN_POINTER(NULL);
diff --git a/src/test/regress/expected/jsonb.out b/src/test/regress/expected/jsonb.out
index b96c095525c..3198442352c 100644
--- a/src/test/regress/expected/jsonb.out
+++ b/src/test/regress/expected/jsonb.out
@@ -1474,6 +1474,163 @@ SELECT jsonb_extract_path_float4('{"a":3.14}'::jsonb, ARRAY['a']);
                       3.14
 (1 row)
 
+-- Optimized typed extraction: multi-subscript chains
+-- The planner lowers j['a']['b'], j['a'][0], etc. to the extract-path typed
+-- extractor family, reusing the same functions as the #> operator path.
+-- Section M1: planner rewrite verification for multi-subscript chains
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field6']['f1'])::int4 FROM test_jsonb WHERE json_type = 'object';
+                                   QUERY PLAN                                    
+---------------------------------------------------------------------------------
+ Seq Scan on pg_temp.test_jsonb
+   Output: jsonb_extract_path_int4(test_json, ARRAY['field6'::text, 'f1'::text])
+   Filter: (test_jsonb.json_type = 'object'::text)
+(3 rows)
+
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field5'][0])::int4 FROM test_jsonb WHERE json_type = 'object';
+                                   QUERY PLAN                                   
+--------------------------------------------------------------------------------
+ Seq Scan on pg_temp.test_jsonb
+   Output: jsonb_extract_path_int4(test_json, ARRAY['field5'::text, '0'::text])
+   Filter: (test_jsonb.json_type = 'object'::text)
+(3 rows)
+
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_arr[6]['k'])::int4 FROM test_jsonb_arr;
+                                QUERY PLAN                                
+--------------------------------------------------------------------------
+ Seq Scan on pg_temp.test_jsonb_arr
+   Output: jsonb_extract_path_int4(test_arr, ARRAY['6'::text, 'k'::text])
+(2 rows)
+
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field6']['f1'])::float8 FROM test_jsonb WHERE json_type = 'object';
+                                    QUERY PLAN                                     
+-----------------------------------------------------------------------------------
+ Seq Scan on pg_temp.test_jsonb
+   Output: jsonb_extract_path_float8(test_json, ARRAY['field6'::text, 'f1'::text])
+   Filter: (test_jsonb.json_type = 'object'::text)
+(3 rows)
+
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field6']['f1'])::numeric FROM test_jsonb WHERE json_type = 'object';
+                                     QUERY PLAN                                     
+------------------------------------------------------------------------------------
+ Seq Scan on pg_temp.test_jsonb
+   Output: jsonb_extract_path_numeric(test_json, ARRAY['field6'::text, 'f1'::text])
+   Filter: (test_jsonb.json_type = 'object'::text)
+(3 rows)
+
+-- Verify single subscript still uses the existing object-field family, not extract-path
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field4'])::int4 FROM test_jsonb WHERE json_type = 'object';
+                          QUERY PLAN                          
+--------------------------------------------------------------
+ Seq Scan on pg_temp.test_jsonb
+   Output: jsonb_object_field_int4(test_json, 'field4'::text)
+   Filter: (test_jsonb.json_type = 'object'::text)
+(3 rows)
+
+-- Section M1b: multi-subscript int2/float4
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field6']['f1'])::int2 FROM test_jsonb WHERE json_type = 'object';
+                                   QUERY PLAN                                    
+---------------------------------------------------------------------------------
+ Seq Scan on pg_temp.test_jsonb
+   Output: jsonb_extract_path_int2(test_json, ARRAY['field6'::text, 'f1'::text])
+   Filter: (test_jsonb.json_type = 'object'::text)
+(3 rows)
+
+SELECT (('{"a":{"b":7}}'::jsonb)['a']['b'])::int2;
+ int2 
+------
+    7
+(1 row)
+
+-- Section M2: execution through multi-subscript rewrite
+SELECT (('{"a":{"b":42}}'::jsonb)['a']['b'])::int4;
+ int4 
+------
+   42
+(1 row)
+
+SELECT (('{"a":[10,20,30]}'::jsonb)['a'][1])::int4;
+ int4 
+------
+   20
+(1 row)
+
+SELECT (('[{"a":true}]'::jsonb)[0]['a'])::bool;
+ bool 
+------
+ t
+(1 row)
+
+SELECT (('[[1,2],[3,4]]'::jsonb)[0][1])::int4;
+ int4 
+------
+    2
+(1 row)
+
+SELECT (('{"a":[{"b":3.14}]}'::jsonb)['a'][0]['b'])::float8;
+ float8 
+--------
+   3.14
+(1 row)
+
+SELECT (('{"a":{"b":42}}'::jsonb)['a']['b'])::numeric;
+ numeric 
+---------
+      42
+(1 row)
+
+SELECT (('{"a":{"b":42}}'::jsonb)['a']['b'])::int8;
+ int8 
+------
+   42
+(1 row)
+
+-- negative index in nested chain
+SELECT (('{"a":[10,20,30]}'::jsonb)['a'][-1])::int4;
+ int4 
+------
+   30
+(1 row)
+
+-- Section M3: NULL semantics for multi-subscript chains
+SELECT (('{"a":1}'::jsonb)['x']['b'])::int4;  -- missing intermediate key
+ int4 
+------
+     
+(1 row)
+
+SELECT (('{"a":{"c":1}}'::jsonb)['a']['b'])::int4;  -- missing final key
+ int4 
+------
+     
+(1 row)
+
+SELECT (('{"a":{"b":null}}'::jsonb)['a']['b'])::int4;  -- JSON null leaf
+ int4 
+------
+     
+(1 row)
+
+SELECT (('{"a":[1]}'::jsonb)['a'][5])::int4;  -- out-of-range nested index
+ int4 
+------
+     
+(1 row)
+
+-- Section M4: type-mismatch errors for multi-subscript chains
+SELECT (('{"a":{"b":"hello"}}'::jsonb)['a']['b'])::int4;  -- string to int4
+ERROR:  cannot cast jsonb string to type integer
+SELECT (('{"a":{"b":[1,2]}}'::jsonb)['a']['b'])::int4;  -- container to int4
+ERROR:  cannot cast jsonb array to type integer
+-- Section M5: non-constant int4 subscript declines rewrite
+EXPLAIN (VERBOSE, COSTS OFF)
+SELECT (('{"a":[[10,20],[30,40]]}'::jsonb)['a'][i][0])::int4 FROM generate_series(0,1) AS i;
+                                  QUERY PLAN                                  
+------------------------------------------------------------------------------
+ Function Scan on pg_catalog.generate_series i
+   Output: (('{"a": [[10, 20], [30, 40]]}'::jsonb)['a'::text][i][0])::integer
+   Function Call: generate_series(0, 1)
+(3 rows)
+
 SELECT test_json -> 'x' FROM test_jsonb WHERE json_type = 'scalar';
  ?column? 
 ----------
diff --git a/src/test/regress/sql/jsonb.sql b/src/test/regress/sql/jsonb.sql
index 11191fd0f82..ed70c4d29ee 100644
--- a/src/test/regress/sql/jsonb.sql
+++ b/src/test/regress/sql/jsonb.sql
@@ -419,6 +419,49 @@ SELECT jsonb_extract_path_int2('{"a":{"b":7}}'::jsonb, ARRAY['a','b']);
 SELECT jsonb_extract_path_float4('{"a":3.14}'::jsonb, ARRAY['a']);
 
 
+-- Optimized typed extraction: multi-subscript chains
+-- The planner lowers j['a']['b'], j['a'][0], etc. to the extract-path typed
+-- extractor family, reusing the same functions as the #> operator path.
+
+-- Section M1: planner rewrite verification for multi-subscript chains
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field6']['f1'])::int4 FROM test_jsonb WHERE json_type = 'object';
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field5'][0])::int4 FROM test_jsonb WHERE json_type = 'object';
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_arr[6]['k'])::int4 FROM test_jsonb_arr;
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field6']['f1'])::float8 FROM test_jsonb WHERE json_type = 'object';
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field6']['f1'])::numeric FROM test_jsonb WHERE json_type = 'object';
+
+-- Verify single subscript still uses the existing object-field family, not extract-path
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field4'])::int4 FROM test_jsonb WHERE json_type = 'object';
+
+
+-- Section M1b: multi-subscript int2/float4
+EXPLAIN (VERBOSE, COSTS OFF) SELECT (test_json['field6']['f1'])::int2 FROM test_jsonb WHERE json_type = 'object';
+SELECT (('{"a":{"b":7}}'::jsonb)['a']['b'])::int2;
+-- Section M2: execution through multi-subscript rewrite
+SELECT (('{"a":{"b":42}}'::jsonb)['a']['b'])::int4;
+SELECT (('{"a":[10,20,30]}'::jsonb)['a'][1])::int4;
+SELECT (('[{"a":true}]'::jsonb)[0]['a'])::bool;
+SELECT (('[[1,2],[3,4]]'::jsonb)[0][1])::int4;
+SELECT (('{"a":[{"b":3.14}]}'::jsonb)['a'][0]['b'])::float8;
+SELECT (('{"a":{"b":42}}'::jsonb)['a']['b'])::numeric;
+SELECT (('{"a":{"b":42}}'::jsonb)['a']['b'])::int8;
+-- negative index in nested chain
+SELECT (('{"a":[10,20,30]}'::jsonb)['a'][-1])::int4;
+
+-- Section M3: NULL semantics for multi-subscript chains
+SELECT (('{"a":1}'::jsonb)['x']['b'])::int4;  -- missing intermediate key
+SELECT (('{"a":{"c":1}}'::jsonb)['a']['b'])::int4;  -- missing final key
+SELECT (('{"a":{"b":null}}'::jsonb)['a']['b'])::int4;  -- JSON null leaf
+SELECT (('{"a":[1]}'::jsonb)['a'][5])::int4;  -- out-of-range nested index
+
+-- Section M4: type-mismatch errors for multi-subscript chains
+SELECT (('{"a":{"b":"hello"}}'::jsonb)['a']['b'])::int4;  -- string to int4
+SELECT (('{"a":{"b":[1,2]}}'::jsonb)['a']['b'])::int4;  -- container to int4
+
+-- Section M5: non-constant int4 subscript declines rewrite
+EXPLAIN (VERBOSE, COSTS OFF)
+SELECT (('{"a":[[10,20],[30,40]]}'::jsonb)['a'][i][0])::int4 FROM generate_series(0,1) AS i;
+
 SELECT test_json -> 'x' FROM test_jsonb WHERE json_type = 'scalar';
 SELECT test_json -> 'x' FROM test_jsonb WHERE json_type = 'array';
 SELECT test_json -> 'x' FROM test_jsonb WHERE json_type = 'object';
-- 
2.52.0