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

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

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

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index 668fe5d358..129699aa8a 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -20426,12 +20426,19 @@ SELECT NULLIF(value, '(none)') ...
         <indexterm>
          <primary>array_sort</primary>
         </indexterm>
-        <function>array_sort</function> ( <type>anyarray</type> <optional> COLLATE <replaceable>collation_name</replaceable> </optional>)
+        <function>array_sort</function> ( <type>anyarray</type> <optional> COLLATE <replaceable>collation_name</replaceable> </optional> <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. If the <literal>COLLATE</literal> option is specified then sorting is based on <replaceable>collation_name</replaceable>, otherwise
+        Sorts the first dimension of the array.
+       </para>
+       <para>
+        If the <literal>COLLATE</literal> option is specified then sorting is based on <replaceable>collation_name</replaceable>, otherwise
         using array element type's collation.
+        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 9b775313ca..1ff61033b2 100644
--- a/src/backend/utils/adt/array_userfuncs.c
+++ b/src/backend/utils/adt/array_userfuncs.c
@@ -1698,7 +1698,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)
@@ -1714,10 +1713,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)
@@ -1735,8 +1747,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",
@@ -1756,8 +1770,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",
@@ -1767,9 +1783,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))
@@ -1795,3 +1811,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 79641da6ff..9bbc97445c 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -1744,6 +1744,12 @@
 { oid => '8810', descr => 'sort array',
   proname => 'array_sort', provolatile => 'v', 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 ba1da71da1..f4781233e0 100644
--- a/src/test/regress/expected/arrays.out
+++ b/src/test/regress/expected/arrays.out
@@ -2728,6 +2728,32 @@ SELECT array_sort('{foo,bar,CCC,Abc,bbc}'::text[] COLLATE "C");
  {Abc,CCC,bar,bbc,foo}
 (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[], 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 b2a814d195..b98e111329 100644
--- a/src/test/regress/sql/arrays.sql
+++ b/src/test/regress/sql/arrays.sql
@@ -834,5 +834,13 @@ 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[]);
 SELECT array_sort('{foo,bar,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[], false, false);
+
 -- multidimensional array tests
 SELECT array_sort(ARRAY[[2,4],[2,1],[6,5]]);
-- 
2.39.5