pg_experiments, Part 1: Modifying operator logic

Exploring Postgres internals by changing operator logic.

Disclaimer: This is my first time diving deep into Postgres source code. I wrote this post after my short, 1-hour exploration in the code base, so I'm unaware of any ideal tricks/techniques to use for implementing things.

This is my first microblog! I’ll be writing short-form content for smaller projects/experiments. Edit: Moved this to the main blog.

Today’s experiment was something which I thought would be fun to write. Thanks to Phil Eaton for helping me out with this experiment!

Why?

During a conversation, Phil wrote this:

I’ve gone looking for where builtin SQL functions like +, - are defined and got lost. So if someone wants to find that and share that would be nice! 🙂

I got curious and wanted to know more. And since the job market has been meh this year, I decided to explore this further to overcome the monotonic, boring job search.

The goal for this exploration is to find where this logic is present, and tweak it to make the + operator subtract the passed values.

Diving deep into source code

Since Phil mentioned about the + operator, I thought it would be a good start to start checking references for “add” as a keyword for the search.

I came across numeric.c in src/backend/utils/adt, which happens to be a gigantic ~10k+ LOC file. On doing a few more searches, I found something which does addition: numeric_add:

https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l2842

/*
 * numeric_add() -
 *
 *	Add two numerics
 */
Datum
numeric_add(PG_FUNCTION_ARGS)
{
	Numeric		num1 = PG_GETARG_NUMERIC(0);
	Numeric		num2 = PG_GETARG_NUMERIC(1);
	Numeric		res;

	res = numeric_add_opt_error(num1, num2, NULL);

	PG_RETURN_NUMERIC(res);
}

/*
 * numeric_add_opt_error() -
 *
 *	Internal version of numeric_add().  If "*have_error" flag is provided,
 *	on error it's set to true, NULL returned.  This is helpful when caller
 *	need to handle errors by itself.
 */
Numeric
numeric_add_opt_error(Numeric num1, Numeric num2, bool *have_error)
{
	NumericVar	arg1;
	NumericVar	arg2;
	NumericVar	result;
	Numeric		res;

	/*
	 * Handle NaN and infinities
	 */
	if (NUMERIC_IS_SPECIAL(num1) || NUMERIC_IS_SPECIAL(num2))
	{
		if (NUMERIC_IS_NAN(num1) || NUMERIC_IS_NAN(num2))
			return make_result(&const_nan);
		if (NUMERIC_IS_PINF(num1))
		{
			if (NUMERIC_IS_NINF(num2))
				return make_result(&const_nan); /* Inf + -Inf */
			else
				return make_result(&const_pinf);
		}
		if (NUMERIC_IS_NINF(num1))
		{
			if (NUMERIC_IS_PINF(num2))
				return make_result(&const_nan); /* -Inf + Inf */
			else
				return make_result(&const_ninf);
		}
		/* by here, num1 must be finite, so num2 is not */
		if (NUMERIC_IS_PINF(num2))
			return make_result(&const_pinf);
		Assert(NUMERIC_IS_NINF(num2));
		return make_result(&const_ninf);
	}

	/*
	 * Unpack the values, let add_var() compute the result and return it.
	 */
	init_var_from_num(num1, &arg1);
	init_var_from_num(num2, &arg2);

	init_var(&result);
	add_var(&arg1, &arg2, &result);

	res = make_result_opt_error(&result, have_error);

	free_var(&result);

	return res;
}

numeric_add_opt_error calls add_var, which has sign handling logic, and calls other functions for addition/subtraction based on signs:

https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l8353

/*
 * add_var() -
 *
 *	Full version of add functionality on variable level (handling signs).
 *	result might point to one of the operands too without danger.
 */
static void
add_var(const NumericVar *var1, const NumericVar *var2, NumericVar *result)
{
	/*
	 * Decide on the signs of the two variables what to do
	 */
	if (var1->sign == NUMERIC_POS)
	{
		if (var2->sign == NUMERIC_POS)
		{
			/*
			 * Both are positive result = +(ABS(var1) + ABS(var2))
			 */
			add_abs(var1, var2, result);
			result->sign = NUMERIC_POS;
		}
		else
		{
			/*
			 * var1 is positive, var2 is negative Must compare absolute values
			 */
			switch (cmp_abs(var1, var2))
			{
				case 0:
					/* ----------
					 * ABS(var1) == ABS(var2)
					 * result = ZERO
					 * ----------
					 */
					zero_var(result);
					result->dscale = Max(var1->dscale, var2->dscale);
					break;

				case 1:
					/* ----------
					 * ABS(var1) > ABS(var2)
					 * result = +(ABS(var1) - ABS(var2))
					 * ----------
					 */
					sub_abs(var1, var2, result);
					result->sign = NUMERIC_POS;
					break;

				case -1:
					/* ----------
					 * ABS(var1) < ABS(var2)
					 * result = -(ABS(var2) - ABS(var1))
					 * ----------
					 */
					sub_abs(var2, var1, result);
					result->sign = NUMERIC_NEG;
					break;
			}
		}
	}
	else
	{
		if (var2->sign == NUMERIC_POS)
		{
			/* ----------
			 * var1 is negative, var2 is positive
			 * Must compare absolute values
			 * ----------
			 */
			switch (cmp_abs(var1, var2))
			{
				case 0:
					/* ----------
					 * ABS(var1) == ABS(var2)
					 * result = ZERO
					 * ----------
					 */
					zero_var(result);
					result->dscale = Max(var1->dscale, var2->dscale);
					break;

				case 1:
					/* ----------
					 * ABS(var1) > ABS(var2)
					 * result = -(ABS(var1) - ABS(var2))
					 * ----------
					 */
					sub_abs(var1, var2, result);
					result->sign = NUMERIC_NEG;
					break;

				case -1:
					/* ----------
					 * ABS(var1) < ABS(var2)
					 * result = +(ABS(var2) - ABS(var1))
					 * ----------
					 */
					sub_abs(var2, var1, result);
					result->sign = NUMERIC_POS;
					break;
			}
		}
		else
		{
			/* ----------
			 * Both are negative
			 * result = -(ABS(var1) + ABS(var2))
			 * ----------
			 */
			add_abs(var1, var2, result);
			result->sign = NUMERIC_NEG;
		}
	}
}

As we saw above, add_var uses the add_abs and sub_abs functions based on the signs. I took a look at add_abs, and it seems to be the exact place where the actual logic happens (carry bits, etc.):

https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l11331

/*
 * add_abs() -
 *
 *	Add the absolute values of two variables into result.
 *	result might point to one of the operands without danger.
 */
static void
add_abs(const NumericVar *var1, const NumericVar *var2, NumericVar *result)
{
	NumericDigit *res_buf;
	NumericDigit *res_digits;
	int			res_ndigits;
	int			res_weight;
	int			res_rscale,
				rscale1,
				rscale2;
	int			res_dscale;
	int			i,
				i1,
				i2;
	int			carry = 0;

	/* copy these values into local vars for speed in inner loop */
	int			var1ndigits = var1->ndigits;
	int			var2ndigits = var2->ndigits;
	NumericDigit *var1digits = var1->digits;
	NumericDigit *var2digits = var2->digits;

	res_weight = Max(var1->weight, var2->weight) + 1;

	res_dscale = Max(var1->dscale, var2->dscale);

	/* Note: here we are figuring rscale in base-NBASE digits */
	rscale1 = var1->ndigits - var1->weight - 1;
	rscale2 = var2->ndigits - var2->weight - 1;
	res_rscale = Max(rscale1, rscale2);

	res_ndigits = res_rscale + res_weight + 1;
	if (res_ndigits <= 0)
		res_ndigits = 1;

	res_buf = digitbuf_alloc(res_ndigits + 1);
	res_buf[0] = 0;				/* spare digit for later rounding */
	res_digits = res_buf + 1;

	i1 = res_rscale + var1->weight + 1;
	i2 = res_rscale + var2->weight + 1;
	for (i = res_ndigits - 1; i >= 0; i--)
	{
		i1--;
		i2--;
		if (i1 >= 0 && i1 < var1ndigits)
			carry += var1digits[i1];
		if (i2 >= 0 && i2 < var2ndigits)
			carry += var2digits[i2];

		if (carry >= NBASE)
		{
			res_digits[i] = carry - NBASE;
			carry = 1;
		}
		else
		{
			res_digits[i] = carry;
			carry = 0;
		}
	}

	Assert(carry == 0);			/* else we failed to allow for carry out */

	digitbuf_free(result->buf);
	result->ndigits = res_ndigits;
	result->buf = res_buf;
	result->digits = res_digits;
	result->weight = res_weight;
	result->dscale = res_dscale;

	/* Remove leading/trailing zeroes */
	strip_var(result);
}

A simple explanation of how it works:

The result is stored in a NumericVar type, which has a few fields to store the values and some metadata (signs, etc.). The documentation is well-written:

https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l254

typedef struct NumericVar
{
	int			ndigits;		/* # of digits in digits[] - can be 0! */
	int			weight;			/* weight of first digit */
	int			sign;			/* NUMERIC_POS, _NEG, _NAN, _PINF, or _NINF */
	int			dscale;			/* display scale */
	NumericDigit *buf;			/* start of palloc'd space for digits[] */
	NumericDigit *digits;		/* base-NBASE digits */
} NumericVar;

Here’s an interesting comment talking about the format and NBASE:

(https://git.postgresql.org/gitweb/?p=postgresql.git;a=blob;f=src/backend/utils/adt/numeric.c;h=bf61fd7dbc09f6ede5799b0eaa00c306e31a1aa1;hb=HEAD#l55)

 * Numeric values are represented in a base-NBASE floating point format.
 * Each "digit" ranges from 0 to NBASE-1.  The type NumericDigit is signed
 * and wide enough to store a digit.  We assume that NBASE*NBASE can fit in
 * an int.  Although the purely calculational routines could handle any even
 * NBASE that's less than sqrt(INT_MAX), in practice we are only interested
 * in NBASE a power of ten, so that I/O conversions and decimal rounding
 * are easy.  Also, it's actually more efficient if NBASE is rather less than
 * sqrt(INT_MAX), so that there is "headroom" for mul_var and div_var_fast to
 * postpone processing carries.
 *
 * Values of NBASE other than 10000 are considered of historical interest only
 * and are no longer supported in any sense; no mechanism exists for the client
 * to discover the base, so every client supporting binary mode expects the
 * base-10000 format.  If you plan to change this, also note the numeric
 * abbreviation code, which assumes NBASE=10000.
 * ----------
 */

Making changes

Even before making any changes, I had to set up a working environment to compile and test changes. Thankfully, Phil had posted this set of commands with comments for getting started with a simple setup:

$ ./configure --prefix=$(pwd)/test-install
$ make -j8 && make install
$ ./test-install/bin/initdb $(pwd)/test-data # Create a database
$ echo "port=8080" > $(pwd)/postgres.conf # This is the minimal postgres run config
$ mkdir test-run # Place for postgres to put lock files
$ ./test-install/bin/postgres --config-file=$(pwd)/postgres.conf -D $(pwd)/test-data -k $(pwd)/test-run

Now that I had some context on how numeric addition works in Postgres, making the change was really easy; all I had to do was change the call from add_var to sub_var in numeric_add_opt_error. (keep in mind that I’m ignoring other cases)

Diff:

diff --git a/src/backend/utils/adt/numeric.c b/src/backend/utils/adt/numeric.c
index bf61fd7dbc..e174f9a9ff 100644
--- a/src/backend/utils/adt/numeric.c
+++ b/src/backend/utils/adt/numeric.c
@@ -2905,7 +2905,7 @@ numeric_add_opt_error(Numeric num1, Numeric num2, bool *have_error)
        init_var_from_num(num2, &arg2);

        init_var(&result);
-       add_var(&arg1, &arg2, &result);
+       sub_var(&arg1, &arg2, &result);

        res = make_result_opt_error(&result, have_error);

Let’s recompile and restart the server:

$ make -j8 && make install
$ ./test-install/bin/postgres --config-file=$(pwd)/postgres.conf -D $(pwd)/test-data -k $(pwd)/test-run

And run the shell:

$ ./test-install/bin/psql -p 8080 -h localhost -d postgres

To test it out, I created a new table (to prevent possible constant folding and just to be extra sure that it works):

postgres=# CREATE TABLE numbers (num1 numeric, num2 numeric);
CREATE TABLE
postgres=# INSERT INTO numbers (num1, num2) VALUES (5.25, 3.75);
INSERT 0 1

And the moment I had been waiting for:

postgres=# SELECT num1 + num2 AS sum FROM numbers;
 sum
------
 1.50
(1 row)

Yay! 5.25 - 3.75 = 1.50, so it’s working as intended!

What’s next?

The above exploration was something I managed to do in an hour; I’m currently looking to explore more and have a project planned where I’d try to exploit Postgres internals and build something out of it!

Please email (contact <at> aadhav.me) or DM on Twitter if you have any questions, corrections, or if you are hiring.