From d7c2956893222bb025822810860a01539ca3ec85 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Tue, 5 Nov 2024 14:53:42 +0800
Subject: [PATCH v12 2/3] array_sort preserve array dimenion and bound info

---
 src/backend/utils/adt/array_userfuncs.c | 30 +++++++++++++--
 src/backend/utils/adt/arrayfuncs.c      |  3 +-
 src/include/utils/array.h               |  1 +
 src/test/regress/expected/arrays.out    | 49 +++++++++++++++++++++++++
 src/test/regress/sql/arrays.sql         | 18 ++++++++-
 5 files changed, 96 insertions(+), 5 deletions(-)

diff --git a/src/backend/utils/adt/array_userfuncs.c b/src/backend/utils/adt/array_userfuncs.c
index a45586c892..436f761614 100644
--- a/src/backend/utils/adt/array_userfuncs.c
+++ b/src/backend/utils/adt/array_userfuncs.c
@@ -1826,8 +1826,18 @@ array_sort(PG_FUNCTION_ARGS)
 	Datum		value;
 	bool		isnull;
 	ArrayBuildStateAny *astate = NULL;
+	int			ndim,
+			   *dims,
+			   *lbs;
 
-	if (ARR_NDIM(array) < 1)
+	ndim = ARR_NDIM(array);
+	dims = ARR_DIMS(array);
+	lbs = ARR_LBOUND(array);
+
+	if (ndim < 1)
+		PG_RETURN_ARRAYTYPE_P(array);
+
+	if(dims[0] < 2)
 		PG_RETURN_ARRAYTYPE_P(array);
 
 	elmtyp = ARR_ELEMTYPE(array);
@@ -1841,7 +1851,7 @@ array_sort(PG_FUNCTION_ARGS)
 		fcinfo->flinfo->fn_extra = (void *) cache_info;
 	}
 
-	if (ARR_NDIM(array) == 1)
+	if (ndim == 1)
 	{
 		/* Finds the ordering operator for the type for 1-D arrays */
 		typentry = cache_info->typentry;
@@ -1883,7 +1893,7 @@ array_sort(PG_FUNCTION_ARGS)
 										   collation,
 										   false, work_mem, NULL, false);
 
-	array_iterator = array_create_iterator(array, ARR_NDIM(array) - 1, NULL);
+	array_iterator = array_create_iterator(array, ndim - 1, NULL);
 	while (array_iterate(array_iterator, &value, &isnull))
 	{
 		tuplesort_putdatum(tuplesortstate, value, isnull);
@@ -1903,6 +1913,20 @@ array_sort(PG_FUNCTION_ARGS)
 
 	tuplesort_end(tuplesortstate);
 
+	if (astate->arraystate != NULL)
+	{
+		memcpy(astate->arraystate->dims, dims, ndim * sizeof(int));
+		memcpy(astate->arraystate->lbs, lbs, ndim * sizeof(int));
+		Assert(ndim == astate->arraystate->ndims);
+	}
+	else
+	{
+		astate->scalarstate->lbs = lbs[0];
+
+		/* zero dimension won't reach here */
+		Assert(ndim == 1);
+	}
+
 	/* Avoid leaking memory when handed toasted input */
 	PG_FREE_IF_COPY(array, 0);
 	PG_RETURN_DATUM(makeArrayResultAny(astate, CurrentMemoryContext, true));
diff --git a/src/backend/utils/adt/arrayfuncs.c b/src/backend/utils/adt/arrayfuncs.c
index a715e7e0b8..ffdf3449e9 100644
--- a/src/backend/utils/adt/arrayfuncs.c
+++ b/src/backend/utils/adt/arrayfuncs.c
@@ -5335,6 +5335,7 @@ initArrayResultWithSize(Oid element_type, MemoryContext rcontext,
 						 &astate->typbyval,
 						 &astate->typalign);
 
+	astate->lbs = 1;	/* default lower-bound value set to 1, see array_in also */
 	return astate;
 }
 
@@ -5867,7 +5868,7 @@ makeArrayResultAny(ArrayBuildStateAny *astate,
 		/* If no elements were presented, we want to create an empty array */
 		ndims = (astate->scalarstate->nelems > 0) ? 1 : 0;
 		dims[0] = astate->scalarstate->nelems;
-		lbs[0] = 1;
+		lbs[0] = astate->scalarstate->lbs;
 
 		result = makeMdArrayResult(astate->scalarstate, ndims, dims, lbs,
 								   rcontext, release);
diff --git a/src/include/utils/array.h b/src/include/utils/array.h
index 157cc0e4c6..7b7937d494 100644
--- a/src/include/utils/array.h
+++ b/src/include/utils/array.h
@@ -193,6 +193,7 @@ typedef struct ArrayBuildState
 	int			nelems;			/* number of valid entries in above arrays */
 	Oid			element_type;	/* data type of the Datums */
 	int16		typlen;			/* needed info about datatype */
+	int			lbs;			/* one dimension, one lower-bound is OK*/
 	bool		typbyval;
 	char		typalign;
 	bool		private_cxt;	/* use private memory context */
diff --git a/src/test/regress/expected/arrays.out b/src/test/regress/expected/arrays.out
index 846d6acd85..9e371902d3 100644
--- a/src/test/regress/expected/arrays.out
+++ b/src/test/regress/expected/arrays.out
@@ -2735,6 +2735,18 @@ SELECT array_reverse('{{1,2},{3,4},{5,6},{7,8}}'::int[]);
 (1 row)
 
 -- array_sort
+SELECT array_sort('{1}'::int[]);
+ array_sort 
+------------
+ {1}
+(1 row)
+
+SELECT array_sort('{{2,1}}'::int[]);
+ array_sort 
+------------
+ {{2,1}}
+(1 row)
+
 SELECT array_sort('{}'::int[]);
  array_sort 
 ------------
@@ -2772,3 +2784,40 @@ SELECT array_sort(ARRAY[[2,4],[2,1],[6,5]]);
  {{2,1},{2,4},{6,5}}
 (1 row)
 
+create table t(a int[]);
+insert into t values ('[-1:-0]={7,1}'::int[]),
+('[-2:-0][20:21]={{1,2},{10,20},{1,-4}}'),
+('[-2:-0][20:22]={{-11,2,-1},{-11,2, 1},{-11,-4, 10}}'),
+('[-13:-10][0:1][20:22]={
+{{1,2,112},{1,2,-123}},
+{{10,-20,1},{11,123,3}},
+{{10,-20,1},{11,-123,-9}},
+{{1,2,-11},{1,2,211}}}'::int[]);
+SELECT array_sort(t.a) from t;
+                                                       array_sort                                                       
+------------------------------------------------------------------------------------------------------------------------
+ [-1:0]={1,7}
+ [-2:0][20:21]={{1,-4},{1,2},{10,20}}
+ [-2:0][20:22]={{-11,-4,10},{-11,2,-1},{-11,2,1}}
+ [-13:-10][0:1][20:22]={{{1,2,-11},{1,2,211}},{{1,2,112},{1,2,-123}},{{10,-20,1},{11,-123,-9}},{{10,-20,1},{11,123,3}}}
+(4 rows)
+
+SELECT array_sort((t.a) [-13:-10][0:1][21:22]) from t where array_ndims(a) = 3;
+                                  array_sort                                  
+------------------------------------------------------------------------------
+ {{{-20,1},{-123,-9}},{{-20,1},{123,3}},{{2,-11},{2,211}},{{2,112},{2,-123}}}
+(1 row)
+
+SELECT array_sort((t.a) [-13:-11][0:1][21:22]) from t where array_ndims(a) = 3;
+                         array_sort                         
+------------------------------------------------------------
+ {{{-20,1},{-123,-9}},{{-20,1},{123,3}},{{2,112},{2,-123}}}
+(1 row)
+
+SELECT array_sort((t.a) [-13:-11][0:0][20:21]) from t where array_ndims(a) = 3;
+           array_sort            
+---------------------------------
+ {{{1,2}},{{10,-20}},{{10,-20}}}
+(1 row)
+
+drop table t;
diff --git a/src/test/regress/sql/arrays.sql b/src/test/regress/sql/arrays.sql
index 17205f483b..2a40c401b5 100644
--- a/src/test/regress/sql/arrays.sql
+++ b/src/test/regress/sql/arrays.sql
@@ -836,6 +836,8 @@ SELECT array_reverse('{1,2,3,NULL,4,5,6}'::int[]);
 SELECT array_reverse('{{1,2},{3,4},{5,6},{7,8}}'::int[]);
 
 -- array_sort
+SELECT array_sort('{1}'::int[]);
+SELECT array_sort('{{2,1}}'::int[]);
 SELECT array_sort('{}'::int[]);
 SELECT array_sort('{1,3,5,2,4,6}'::int[]);
 SELECT array_sort('{1.1,3.3,5.5,2.2,4.4,6.6}'::numeric[]);
@@ -843,4 +845,18 @@ SELECT array_sort('{foo,bar,CCC,Abc,bbc}'::text[] COLLATE "C");
 SELECT array_sort('{foo,bar,null,CCC,Abc,bbc}'::text[] COLLATE "C");
 
 -- multidimensional array tests
-SELECT array_sort(ARRAY[[2,4],[2,1],[6,5]]);
\ No newline at end of file
+SELECT array_sort(ARRAY[[2,4],[2,1],[6,5]]);
+create table t(a int[]);
+insert into t values ('[-1:-0]={7,1}'::int[]),
+('[-2:-0][20:21]={{1,2},{10,20},{1,-4}}'),
+('[-2:-0][20:22]={{-11,2,-1},{-11,2, 1},{-11,-4, 10}}'),
+('[-13:-10][0:1][20:22]={
+{{1,2,112},{1,2,-123}},
+{{10,-20,1},{11,123,3}},
+{{10,-20,1},{11,-123,-9}},
+{{1,2,-11},{1,2,211}}}'::int[]);
+SELECT array_sort(t.a) from t;
+SELECT array_sort((t.a) [-13:-10][0:1][21:22]) from t where array_ndims(a) = 3;
+SELECT array_sort((t.a) [-13:-11][0:1][21:22]) from t where array_ndims(a) = 3;
+SELECT array_sort((t.a) [-13:-11][0:0][20:21]) from t where array_ndims(a) = 3;
+drop table t;
\ No newline at end of file
-- 
2.34.1

