From dcd8d776feaf8c295b51fa2544846c6173ee664c Mon Sep 17 00:00:00 2001
From: Andy Fan <zhihui.fan1213@gmail.com>
Date: Mon, 7 Aug 2023 10:44:23 +0800
Subject: [PATCH v3] Optimize extracting a given data type from jsonb.

Previously after we get a JsonbValue, we need to convert it to
Jsonb first then cast the Jsonb to the given type. In this patch,
we covert the JsonbValue to the desired type directly.
---
 src/backend/utils/adt/jsonb.c            | 134 +++++++++++++++++++++++
 src/include/catalog/catversion.h         |   2 +-
 src/include/catalog/pg_proc.dat          |  22 ++--
 src/test/regress/expected/jsonb.out      | 134 ++++++++++++-----------
 src/test/regress/expected/opr_sanity.out |   7 +-
 src/test/regress/sql/jsonb.sql           |  36 +++---
 6 files changed, 244 insertions(+), 91 deletions(-)

diff --git a/src/backend/utils/adt/jsonb.c b/src/backend/utils/adt/jsonb.c
index 9781852b0cb..a577b2a279f 100644
--- a/src/backend/utils/adt/jsonb.c
+++ b/src/backend/utils/adt/jsonb.c
@@ -17,11 +17,14 @@
 #include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
 #include "funcapi.h"
+#include "nodes/makefuncs.h"
+#include "nodes/supportnodes.h"
 #include "libpq/pqformat.h"
 #include "miscadmin.h"
 #include "utils/builtins.h"
 #include "utils/date.h"
 #include "utils/datetime.h"
+#include "utils/fmgroids.h"
 #include "utils/json.h"
 #include "utils/jsonb.h"
 #include "utils/jsonfuncs.h"
@@ -2038,6 +2041,137 @@ cannotCastJsonbValue(enum jbvType type, const char *sqltype)
 	elog(ERROR, "unknown jsonb type: %d", (int) type);
 }
 
+static bool
+jsonb_cast_is_optimized(Oid target_type)
+{
+	switch(target_type)
+	{
+		case NUMERICOID:
+		case BOOLOID:
+		case INT2OID:
+		case INT4OID:
+		case INT8OID:
+		case FLOAT4OID:
+		case FLOAT8OID:
+			return true;
+		default:
+			return false;
+	}
+}
+
+Datum
+jsonb_cast_support(PG_FUNCTION_ARGS)
+{
+	Node	   *rawreq = (Node *) PG_GETARG_POINTER(0);
+
+	if (IsA(rawreq, SupportRequestSimplify))
+	{
+		SupportRequestSimplify *req = (SupportRequestSimplify *) rawreq;
+		FuncExpr	*fexpr = palloc(sizeof(FuncExpr));
+		OpExpr		*opexpr;
+
+		memcpy(fexpr, req->fcall, sizeof(FuncExpr));
+
+		opexpr = (OpExpr *) linitial(fexpr->args);
+
+		/*
+		 * Simplify cast(jsonb_object_filed(jsonb, filedName) as type)
+		 * to jsonb_object_field_type(jsonb, filedName, targetTypeOid);
+		 */
+		if (IsA(opexpr, OpExpr) &&
+			opexpr->opfuncid  == F_JSONB_OBJECT_FIELD &&
+			jsonb_cast_is_optimized(fexpr->funcresulttype))
+		{
+			fexpr->funcid = F_JSONB_OBJECT_FIELD_TYPE;
+			fexpr->args = opexpr->args;
+
+			/* Tell the jsonb_object_field_type what is the target type. */
+			fexpr->args = lappend(fexpr->args, makeConst(OIDOID, 0, 0, sizeof(Oid),
+														 fexpr->funcresulttype,
+														 false, true));
+		}
+
+		PG_RETURN_POINTER(fexpr);
+	}
+
+	PG_RETURN_POINTER(NULL);
+}
+
+Datum
+jsonb_object_field_type(PG_FUNCTION_ARGS)
+{
+	Jsonb	   *jb = PG_GETARG_JSONB_P(0);
+	text	   *key = PG_GETARG_TEXT_PP(1);
+	Oid			targetOid = PG_GETARG_OID(2);
+
+	JsonbValue *v;
+	JsonbValue	vbuf;
+
+	if (!JB_ROOT_IS_OBJECT(jb))
+		PG_RETURN_NULL();
+
+	v = getKeyJsonValueFromContainer(&jb->root,
+									 VARDATA_ANY(key),
+									 VARSIZE_ANY_EXHDR(key),
+									 &vbuf);
+
+	if (v == NULL)
+		PG_RETURN_NULL();
+
+	switch(targetOid)
+	{
+		Datum	retValue;
+
+		case BOOLOID:
+			if (v->type != jbvBool)
+				cannotCastJsonbValue(v->type, "bool");
+			PG_RETURN_BOOL(v->val.boolean);
+
+		case NUMERICOID:
+			if (v->type != jbvNumeric)
+				cannotCastJsonbValue(v->type, "numeric");
+			PG_RETURN_NUMERIC(v->val.numeric);
+		case INT2OID:
+			if (v->type != jbvNumeric)
+				cannotCastJsonbValue(v->type, "smallint");
+			retValue = DirectFunctionCall1(numeric_int2,
+										   NumericGetDatum(v->val.numeric));
+			PG_RETURN_DATUM(retValue);
+		case INT4OID:
+			if (v->type != jbvNumeric)
+				cannotCastJsonbValue(v->type, "integer");
+			retValue = DirectFunctionCall1(numeric_int4,
+										   NumericGetDatum(v->val.numeric));
+			PG_RETURN_DATUM(retValue);
+
+		case INT8OID:
+			if (v->type != jbvNumeric)
+				cannotCastJsonbValue(v->type, "bigint");
+			retValue = DirectFunctionCall1(numeric_int8,
+										   NumericGetDatum(v->val.numeric));
+			PG_RETURN_DATUM(retValue);
+
+		case FLOAT4OID:
+			if (v->type != jbvNumeric)
+				cannotCastJsonbValue(v->type, "real");
+			retValue = DirectFunctionCall1(numeric_float4,
+										   NumericGetDatum(v->val.numeric));
+			PG_RETURN_DATUM(retValue);
+
+		case FLOAT8OID:
+			if (v->type != jbvNumeric)
+				cannotCastJsonbValue(v->type, "double precision");
+			retValue = DirectFunctionCall1(numeric_float8,
+										   NumericGetDatum(v->val.numeric));
+			PG_RETURN_DATUM(retValue);
+
+		default:
+			elog(ERROR, "cast jsonb field to %d is not supported.", targetOid);
+	}
+
+	PG_RETURN_POINTER(0);
+}
+
 Datum
 jsonb_bool(PG_FUNCTION_ARGS)
 {
diff --git a/src/include/catalog/catversion.h b/src/include/catalog/catversion.h
index f507b49bb28..391dfb81b2a 100644
--- a/src/include/catalog/catversion.h
+++ b/src/include/catalog/catversion.h
@@ -57,6 +57,6 @@
  */
 
 /*							yyyymmddN */
-#define CATALOG_VERSION_NO	202307261
+#define CATALOG_VERSION_NO	202308071
 
 #endif
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 6996073989a..108280b355d 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -4575,25 +4575,26 @@
   proname => 'pg_lsn', prorettype => 'pg_lsn', proargtypes => 'numeric',
   prosrc => 'numeric_pg_lsn' },
 
-{ oid => '3556', descr => 'convert jsonb to boolean',
+{ oid => '3556', descr => 'convert jsonb to boolean', prosupport => 'jsonb_cast_support',
   proname => 'bool', prorettype => 'bool', proargtypes => 'jsonb',
   prosrc => 'jsonb_bool' },
 { oid => '3449', descr => 'convert jsonb to numeric',
-  proname => 'numeric', prorettype => 'numeric', proargtypes => 'jsonb',
+  proname => 'numeric', prosupport => 'jsonb_cast_support',
+  prorettype => 'numeric', proargtypes => 'jsonb',
   prosrc => 'jsonb_numeric' },
-{ oid => '3450', descr => 'convert jsonb to int2',
+{ oid => '3450', descr => 'convert jsonb to int2', prosupport => 'jsonb_cast_support',
   proname => 'int2', prorettype => 'int2', proargtypes => 'jsonb',
   prosrc => 'jsonb_int2' },
-{ oid => '3451', descr => 'convert jsonb to int4',
+{ oid => '3451', descr => 'convert jsonb to int4', prosupport => 'jsonb_cast_support',
   proname => 'int4', prorettype => 'int4', proargtypes => 'jsonb',
   prosrc => 'jsonb_int4' },
-{ oid => '3452', descr => 'convert jsonb to int8',
+{ oid => '3452', descr => 'convert jsonb to int8', prosupport => 'jsonb_cast_support',
   proname => 'int8', prorettype => 'int8', proargtypes => 'jsonb',
   prosrc => 'jsonb_int8' },
-{ oid => '3453', descr => 'convert jsonb to float4',
+{ oid => '3453', descr => 'convert jsonb to float4', prosupport => 'jsonb_cast_support',
   proname => 'float4', prorettype => 'float4', proargtypes => 'jsonb',
   prosrc => 'jsonb_float4' },
-{ oid => '2580', descr => 'convert jsonb to float8',
+{ oid => '2580', descr => 'convert jsonb to float8', prosupport => 'jsonb_cast_support',
   proname => 'float8', prorettype => 'float8', proargtypes => 'jsonb',
   prosrc => 'jsonb_float8' },
 
@@ -9928,6 +9929,13 @@
   proname => 'jsonb_object_field_text', prorettype => 'text',
   proargtypes => 'jsonb text', proargnames => '{from_json, field_name}',
   prosrc => 'jsonb_object_field_text' },
+{ oid => '3813', descr => 'return a given type specified in desired_type from jsonb field',
+  proname => 'jsonb_object_field_type', prorettype => 'anyelement',
+  proargtypes => 'jsonb text oid', proargnames => '{from_json, field_name, desired_type}',
+  prosrc => 'jsonb_object_field_type'},
+{ oid => '3814', descr => 'planner support for numeric(jsonb)',
+  proname => 'jsonb_cast_support', prorettype => 'internal',
+  proargtypes => 'internal', prosrc => 'jsonb_cast_support' },
 { oid => '3215',
   proname => 'jsonb_array_element', prorettype => 'jsonb',
   proargtypes => 'jsonb int4', proargnames => '{from_json, element_index}',
diff --git a/src/test/regress/expected/jsonb.out b/src/test/regress/expected/jsonb.out
index 4a16d0dbafb..5a144534cc1 100644
--- a/src/test/regress/expected/jsonb.out
+++ b/src/test/regress/expected/jsonb.out
@@ -5471,107 +5471,113 @@ select ts_headline('[]'::jsonb, tsquery('aaa & bbb'));
 (1 row)
 
 -- casts
-select 'true'::jsonb::bool;
- bool 
-------
- t
+select 'true'::jsonb::bool, ('{"a": true}'::jsonb->'a')::bool;
+ bool | bool 
+------+------
+ t    | t
 (1 row)
 
 select '[]'::jsonb::bool;
 ERROR:  cannot cast jsonb array to type boolean
-select '1.0'::jsonb::float;
- float8 
---------
-      1
+select ('{"a": []}'::jsonb->'a')::bool;
+ERROR:  cannot cast jsonb array to type boolean
+select '1.0'::jsonb::float, ('{"a": 1.0}'::jsonb->'a')::float;
+ float8 | float8 
+--------+--------
+      1 |      1
 (1 row)
 
 select '[1.0]'::jsonb::float;
 ERROR:  cannot cast jsonb array to type double precision
-select '12345'::jsonb::int4;
- int4  
--------
- 12345
+select ('{"a": [1.0]}'::jsonb->'a')::float;
+ERROR:  cannot cast jsonb array to type double precision
+select '12345'::jsonb::int4,  ('{"a": 12345}'::jsonb->'a')::int4;
+ int4  | int4  
+-------+-------
+ 12345 | 12345
 (1 row)
 
 select '"hello"'::jsonb::int4;
 ERROR:  cannot cast jsonb string to type integer
-select '12345'::jsonb::numeric;
- numeric 
----------
-   12345
+select ('{"a": "hello"}'::jsonb->'a')::int4;
+ERROR:  cannot cast jsonb string to type integer
+select '12345'::jsonb::numeric, ('{"a": 12345}'::jsonb->'a')::numeric;
+ numeric | numeric 
+---------+---------
+   12345 |   12345
 (1 row)
 
 select '{}'::jsonb::numeric;
 ERROR:  cannot cast jsonb object to type numeric
-select '12345.05'::jsonb::numeric;
- numeric  
-----------
- 12345.05
+select '12345.05'::jsonb::numeric, ('{"a": 12345.05}'::jsonb->'a')::numeric;
+ numeric  | numeric  
+----------+----------
+ 12345.05 | 12345.05
 (1 row)
 
-select '12345.05'::jsonb::float4;
-  float4  
-----------
- 12345.05
+select '12345.05'::jsonb::float4, ('{"a": 12345.05}'::jsonb->'a')::float4;
+  float4  |  float4  
+----------+----------
+ 12345.05 | 12345.05
 (1 row)
 
-select '12345.05'::jsonb::float8;
-  float8  
-----------
- 12345.05
+select '12345.05'::jsonb::float8, ('{"a": 12345.05}'::jsonb->'a')::float8;
+  float8  |  float8  
+----------+----------
+ 12345.05 | 12345.05
 (1 row)
 
-select '12345.05'::jsonb::int2;
- int2  
--------
- 12345
+select '12345.05'::jsonb::int2, ('{"a": 12345.05}'::jsonb->'a')::int2;
+ int2  | int2  
+-------+-------
+ 12345 | 12345
 (1 row)
 
-select '12345.05'::jsonb::int4;
- int4  
--------
- 12345
+select '12345.05'::jsonb::int4, ('{"a": 12345.05}'::jsonb->'a')::int4;
+ int4  | int4  
+-------+-------
+ 12345 | 12345
 (1 row)
 
-select '12345.05'::jsonb::int8;
- int8  
--------
- 12345
+select '12345.05'::jsonb::int8, ('{"a": 12345.05}'::jsonb->'a')::int8;
+ int8  | int8  
+-------+-------
+ 12345 | 12345
 (1 row)
 
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::numeric;
-                       numeric                        
-------------------------------------------------------
- 12345.0000000000000000000000000000000000000000000005
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::numeric, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::numeric;
+                       numeric                        |                       numeric                        
+------------------------------------------------------+------------------------------------------------------
+ 12345.0000000000000000000000000000000000000000000005 | 12345.0000000000000000000000000000000000000000000005
 (1 row)
 
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::float4;
- float4 
---------
-  12345
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::float4,  ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::float4;
+ float4 | float4 
+--------+--------
+  12345 |  12345
 (1 row)
 
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::float8;
- float8 
---------
-  12345
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::float8, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::float8;
+ float8 | float8 
+--------+--------
+  12345 |  12345
 (1 row)
 
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::int2;
- int2  
--------
- 12345
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::int2, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::int2;
+ int2  | int2  
+-------+-------
+ 12345 | 12345
 (1 row)
 
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::int4;
- int4  
--------
- 12345
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::int4, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::int4;
+ int4  | int4  
+-------+-------
+ 12345 | 12345
 (1 row)
 
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::int8;
- int8  
--------
- 12345
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::int8, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::int8;
+ int8  | int8  
+-------+-------
+ 12345 | 12345
 (1 row)
 
diff --git a/src/test/regress/expected/opr_sanity.out b/src/test/regress/expected/opr_sanity.out
index a1bdf2c0b5f..b4c36472eb2 100644
--- a/src/test/regress/expected/opr_sanity.out
+++ b/src/test/regress/expected/opr_sanity.out
@@ -309,8 +309,8 @@ WHERE p1.prorettype IN
      'anyrange'::regtype = ANY (p1.proargtypes) OR
      'anymultirange'::regtype = ANY (p1.proargtypes))
 ORDER BY 2;
- oid  |    proname     
-------+----------------
+ oid  |         proname         
+------+-------------------------
  2296 | anyarray_in
  2502 | anyarray_recv
  2312 | anyelement_in
@@ -320,7 +320,8 @@ ORDER BY 2;
  2400 | array_recv
  3506 | enum_in
  3532 | enum_recv
-(9 rows)
+ 3813 | jsonb_object_field_type
+(10 rows)
 
 -- anyrange and anymultirange are tighter than the rest, can only resolve
 -- from each other
diff --git a/src/test/regress/sql/jsonb.sql b/src/test/regress/sql/jsonb.sql
index e4b7cdf703d..a0b324c8bef 100644
--- a/src/test/regress/sql/jsonb.sql
+++ b/src/test/regress/sql/jsonb.sql
@@ -1496,23 +1496,27 @@ select ts_headline('{}'::jsonb, tsquery('aaa & bbb'));
 select ts_headline('[]'::jsonb, tsquery('aaa & bbb'));
 
 -- casts
-select 'true'::jsonb::bool;
+select 'true'::jsonb::bool, ('{"a": true}'::jsonb->'a')::bool;
 select '[]'::jsonb::bool;
-select '1.0'::jsonb::float;
+select ('{"a": []}'::jsonb->'a')::bool;
+select '1.0'::jsonb::float, ('{"a": 1.0}'::jsonb->'a')::float;
 select '[1.0]'::jsonb::float;
-select '12345'::jsonb::int4;
+select ('{"a": [1.0]}'::jsonb->'a')::float;
+select '12345'::jsonb::int4,  ('{"a": 12345}'::jsonb->'a')::int4;
 select '"hello"'::jsonb::int4;
-select '12345'::jsonb::numeric;
+select ('{"a": "hello"}'::jsonb->'a')::int4;
+
+select '12345'::jsonb::numeric, ('{"a": 12345}'::jsonb->'a')::numeric;
 select '{}'::jsonb::numeric;
-select '12345.05'::jsonb::numeric;
-select '12345.05'::jsonb::float4;
-select '12345.05'::jsonb::float8;
-select '12345.05'::jsonb::int2;
-select '12345.05'::jsonb::int4;
-select '12345.05'::jsonb::int8;
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::numeric;
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::float4;
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::float8;
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::int2;
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::int4;
-select '12345.0000000000000000000000000000000000000000000005'::jsonb::int8;
+select '12345.05'::jsonb::numeric, ('{"a": 12345.05}'::jsonb->'a')::numeric;
+select '12345.05'::jsonb::float4, ('{"a": 12345.05}'::jsonb->'a')::float4;
+select '12345.05'::jsonb::float8, ('{"a": 12345.05}'::jsonb->'a')::float8;
+select '12345.05'::jsonb::int2, ('{"a": 12345.05}'::jsonb->'a')::int2;
+select '12345.05'::jsonb::int4, ('{"a": 12345.05}'::jsonb->'a')::int4;
+select '12345.05'::jsonb::int8, ('{"a": 12345.05}'::jsonb->'a')::int8;
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::numeric, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::numeric;
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::float4,  ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::float4;
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::float8, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::float8;
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::int2, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::int2;
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::int4, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::int4;
+select '12345.0000000000000000000000000000000000000000000005'::jsonb::int8, ('{"a": 12345.0000000000000000000000000000000000000000000005}'::jsonb->'a')::int8;
-- 
2.21.0

