v11-0002-support-sort-order-and-nullsfirst-flag.patch

From 76ed69b5005e8445d68da4da45994748ff83e0ca Mon Sep 17 00:00:00 2001
From: Junwang Zhao <zhjwpku@gmail.com>
Date: Sat, 26 Oct 2024 03:53:17 +0000
Subject: [PATCH v11 2/2] support sort order and nullsfirst flag

Signed-off-by: Junwang Zhao <zhjwpku@gmail.com>
---
 doc/src/sgml/func.sgml                  |  6 +++-
 src/backend/utils/adt/array_userfuncs.c | 42 ++++++++++++++++++++-----
 src/include/catalog/pg_proc.dat         |  6 ++++
 src/test/regress/expected/arrays.out    | 38 ++++++++++++++++++++++
 src/test/regress/sql/arrays.sql         | 10 ++++++
 5 files changed, 94 insertions(+), 8 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index 9e019e18cb..de66991d3d 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -20443,12 +20443,16 @@ SELECT NULLIF(value, '(none)') ...
         <indexterm>
          <primary>array_sort</primary>
         </indexterm>
-        <function>array_sort</function> ( <type>anyarray</type> )
+        <function>array_sort</function> ( <type>anyarray</type> <optional>, <parameter>is_ascending</parameter> <type>boolean</type> <optional>, <parameter>nulls_first</parameter> <type>boolean</type></optional></optional>)
         <returnvalue>anyarray</returnvalue>
        </para>
        <para>
         Sorts the first dimension of the array. The sort order is determined by the <literal>&lt;</literal> operator of the element type, nulls will appear after non-null values.
         The collation to use can be forced by adding a <literal>COLLATE</literal> clause to any of the arguments.
+        If <parameter>is_ascending</parameter> is true then sort by ascending order, otherwise descending order.
+        <parameter>is_ascending</parameter> defaults to true.
+        If <parameter>nulls_first</parameter> is true then nulls appear before non-null values, otherwise nulls appear after non-null values.
+        <parameter>nulls_first</parameter> defaults to the opposite of <parameter>is_ascending</parameter> if not provided.
        </para>
        <para>
         <literal>array_sort(ARRAY[[2,4],[2,1],[6,5]])</literal>
diff --git a/src/backend/utils/adt/array_userfuncs.c b/src/backend/utils/adt/array_userfuncs.c
index a45586c892..df9822d977 100644
--- a/src/backend/utils/adt/array_userfuncs.c
+++ b/src/backend/utils/adt/array_userfuncs.c
@@ -1810,7 +1810,6 @@ typedef struct ArraySortCachedInfo
  * array_sort
  *
  * Sorts the first dimension of the array.
- * The sort order is determined by the "<" operator of the element type.
  */
 Datum
 array_sort(PG_FUNCTION_ARGS)
@@ -1826,10 +1825,23 @@ array_sort(PG_FUNCTION_ARGS)
 	Datum		value;
 	bool		isnull;
 	ArrayBuildStateAny *astate = NULL;
+	bool		is_ascending = true;
+	bool		nulls_first = false;
 
 	if (ARR_NDIM(array) < 1)
 		PG_RETURN_ARRAYTYPE_P(array);
 
+	if (PG_NARGS() > 1)
+	{
+		is_ascending = PG_GETARG_BOOL(1);
+
+		/*
+		 * If nulls_first not provided, it defaults to the opposite of
+		 * is_ascending.
+		 */
+		nulls_first = PG_NARGS() > 2 ? PG_GETARG_BOOL(2) : !is_ascending;
+	}
+
 	elmtyp = ARR_ELEMTYPE(array);
 	cache_info = (ArraySortCachedInfo *) fcinfo->flinfo->fn_extra;
 	if (cache_info == NULL)
@@ -1847,8 +1859,10 @@ array_sort(PG_FUNCTION_ARGS)
 		typentry = cache_info->typentry;
 		if (typentry == NULL || typentry->type_id != elmtyp)
 		{
-			typentry = lookup_type_cache(elmtyp, TYPECACHE_LT_OPR);
-			if (!OidIsValid(typentry->lt_opr))
+			typentry = lookup_type_cache(elmtyp,
+										 is_ascending ? TYPECACHE_LT_OPR : TYPECACHE_GT_OPR);
+			if ((is_ascending && !OidIsValid(typentry->lt_opr)) ||
+				(!is_ascending && !OidIsValid(typentry->gt_opr)))
 				ereport(ERROR,
 						(errcode(ERRCODE_UNDEFINED_FUNCTION),
 						 errmsg("could not identify ordering operator for type %s",
@@ -1868,8 +1882,10 @@ array_sort(PG_FUNCTION_ARGS)
 						(errcode(ERRCODE_UNDEFINED_OBJECT),
 						 errmsg("could not find array type for data type %s",
 								format_type_be(elmtyp))));
-			typentry = lookup_type_cache(array_type, TYPECACHE_LT_OPR);
-			if (!OidIsValid(typentry->lt_opr))
+			typentry = lookup_type_cache(array_type,
+										 is_ascending ? TYPECACHE_LT_OPR : TYPECACHE_GT_OPR);
+			if ((is_ascending && !OidIsValid(typentry->lt_opr)) ||
+				(!is_ascending && !OidIsValid(typentry->gt_opr)))
 				ereport(ERROR,
 						(errcode(ERRCODE_UNDEFINED_FUNCTION),
 						 errmsg("could not identify ordering operator for type %s",
@@ -1879,9 +1895,9 @@ array_sort(PG_FUNCTION_ARGS)
 	}
 
 	tuplesortstate = tuplesort_begin_datum(typentry->type_id,
-										   typentry->lt_opr,
+										   is_ascending ? typentry->lt_opr : typentry->gt_opr,
 										   collation,
-										   false, work_mem, NULL, false);
+										   nulls_first, work_mem, NULL, false);
 
 	array_iterator = array_create_iterator(array, ARR_NDIM(array) - 1, NULL);
 	while (array_iterate(array_iterator, &value, &isnull))
@@ -1907,3 +1923,15 @@ array_sort(PG_FUNCTION_ARGS)
 	PG_FREE_IF_COPY(array, 0);
 	PG_RETURN_DATUM(makeArrayResultAny(astate, CurrentMemoryContext, true));
 }
+
+Datum
+array_sort_order(PG_FUNCTION_ARGS)
+{
+	return array_sort(fcinfo);
+}
+
+Datum
+array_sort_order_nulls_first(PG_FUNCTION_ARGS)
+{
+	return array_sort(fcinfo);
+}
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 55b52e138a..ad6de57e57 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -1747,6 +1747,12 @@
 { oid => '8810', descr => 'sort array',
   proname => 'array_sort', prorettype => 'anyarray',
   proargtypes => 'anyarray', prosrc => 'array_sort'},
+{ oid => '8811', descr => 'sort array',
+  proname => 'array_sort', provolatile => 'v', prorettype => 'anyarray',
+  proargtypes => 'anyarray bool', prosrc => 'array_sort_order'},
+{ oid => '8812', descr => 'sort array',
+  proname => 'array_sort', provolatile => 'v', prorettype => 'anyarray',
+  proargtypes => 'anyarray bool bool', prosrc => 'array_sort_order_nulls_first'},
 { oid => '3816', descr => 'array typanalyze',
   proname => 'array_typanalyze', provolatile => 's', prorettype => 'bool',
   proargtypes => 'internal', prosrc => 'array_typanalyze' },
diff --git a/src/test/regress/expected/arrays.out b/src/test/regress/expected/arrays.out
index 846d6acd85..b0cc38581e 100644
--- a/src/test/regress/expected/arrays.out
+++ b/src/test/regress/expected/arrays.out
@@ -2765,6 +2765,44 @@ SELECT array_sort('{foo,bar,null,CCC,Abc,bbc}'::text[] COLLATE "C");
  {Abc,CCC,bar,bbc,foo,NULL}
 (1 row)
 
+-- array_sort with order specified
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], true);
+           array_sort           
+--------------------------------
+ {1.1,2.2,3.3,4.4,5.5,6.6,NULL}
+(1 row)
+
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], false);
+           array_sort           
+--------------------------------
+ {NULL,6.6,5.5,4.4,3.3,2.2,1.1}
+(1 row)
+
+-- array_sort with order and nullsfirst flag specified
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], true, true);
+           array_sort           
+--------------------------------
+ {NULL,1.1,2.2,3.3,4.4,5.5,6.6}
+(1 row)
+
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], true, false);
+           array_sort           
+--------------------------------
+ {1.1,2.2,3.3,4.4,5.5,6.6,NULL}
+(1 row)
+
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], false, true);
+           array_sort           
+--------------------------------
+ {NULL,6.6,5.5,4.4,3.3,2.2,1.1}
+(1 row)
+
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], false, false);
+           array_sort           
+--------------------------------
+ {6.6,5.5,4.4,3.3,2.2,1.1,NULL}
+(1 row)
+
 -- multidimensional array tests
 SELECT array_sort(ARRAY[[2,4],[2,1],[6,5]]);
      array_sort      
diff --git a/src/test/regress/sql/arrays.sql b/src/test/regress/sql/arrays.sql
index b1f9c97e5c..bc046d9fe4 100644
--- a/src/test/regress/sql/arrays.sql
+++ b/src/test/regress/sql/arrays.sql
@@ -842,5 +842,15 @@ SELECT array_sort('{1.1,3.3,5.5,2.2,4.4,6.6}'::numeric[]);
 SELECT array_sort('{foo,bar,CCC,Abc,bbc}'::text[] COLLATE "C");
 SELECT array_sort('{foo,bar,null,CCC,Abc,bbc}'::text[] COLLATE "C");
 
+-- array_sort with order specified
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], true);
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], false);
+
+-- array_sort with order and nullsfirst flag specified
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], true, true);
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], true, false);
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], false, true);
+SELECT array_sort('{1.1,3.3,5.5,2.2,null,4.4,6.6}'::float8[], false, false);
+
 -- multidimensional array tests
 SELECT array_sort(ARRAY[[2,4],[2,1],[6,5]]);
-- 
2.39.5