feat(hogvm): Added support for ifNull in HogVM (#20549)

Added support for ifNull in HogVM
PostHog · Feb 26, 2024 · c940fa5 · c940fa5
1 parent 4d3b82f
commit c940fa5
Show file tree

Hide file tree

Showing 7 changed files with 39 additions and 9 deletions.
diff --git a/hogvm/README.md b/hogvm/README.md
@@ -85,12 +85,13 @@ execute_bytecode(to_bytecode("'user_id' in cohort 2"), {}, async_operation)
 A PostHog HogQL Bytecode Certified Parser must also implement the following function calls:
 
 ```bash
-concat(...)             # concat('test: ', 1, null, '!') == 'test: 1!'
-match(string, pattern)  # match('fish', '$fi.*') == true
-toString(val)           # toString(true) == 'true'
-toInt(val)              # toInt('123') == 123
-toFloat(val)            # toFloat('123.2') == 123.2
-toUUID(val)             # toUUID('string') == 'string'
+concat(...)              # concat('test: ', 1, null, '!') == 'test: 1!'
+match(string, pattern)   # match('fish', '$fi.*') == true
+toString(val)            # toString(true) == 'true'
+toInt(val)               # toInt('123') == 123
+toFloat(val)             # toFloat('123.2') == 123.2
+toUUID(val)              # toUUID('string') == 'string'
+ifNull(val, alternative) # ifNull('string', false) == 'string'
 ```
 
 ### Null handling

diff --git a/hogvm/python/execute.py b/hogvm/python/execute.py
@@ -130,6 +130,11 @@ def execute_bytecode(bytecode: List[Any], fields: Dict[str, Any]) -> Any:
                             stack.append(int(args[0]) if name == "toInt" else float(args[0]))
                         except ValueError:
                             stack.append(None)
+                    elif name == "ifNull":
+                        if args[0] is not None:
+                            stack.append(args[0])
+                        else:
+                            stack.append(args[1])
                     else:
                         raise HogVMException(f"Unsupported function call: {name}")
                 case _:

diff --git a/hogvm/python/operation.py b/hogvm/python/operation.py
@@ -3,7 +3,7 @@
 HOGQL_BYTECODE_IDENTIFIER = "_h"
 
 
-SUPPORTED_FUNCTIONS = ("concat", "match", "toString", "toInt", "toFloat", "toUUID")
+SUPPORTED_FUNCTIONS = ("concat", "match", "toString", "toInt", "toFloat", "toUUID", "ifNull")
 
 
 class Operation(int, Enum):

diff --git a/hogvm/python/test/test_execute.py b/hogvm/python/test/test_execute.py
@@ -10,7 +10,7 @@
 class TestBytecodeExecute(BaseTest):
     def _run(self, expr: str) -> Any:
         fields = {
-            "properties": {"foo": "bar"},
+            "properties": {"foo": "bar", "nullValue": None},
         }
         return execute_bytecode(create_bytecode(parse_expr(expr)), fields)
 
@@ -51,6 +51,8 @@ def test_bytecode_create(self):
         self.assertEqual(self._run("'a' not in 'car'"), False)
         self.assertEqual(self._run("properties.bla"), None)
         self.assertEqual(self._run("properties.foo"), "bar")
+        self.assertEqual(self._run("ifNull(properties.foo, false)"), "bar")
+        self.assertEqual(self._run("ifNull(properties.nullValue, false)"), False)
         self.assertEqual(self._run("concat('arg', 'another')"), "arganother")
         self.assertEqual(self._run("concat(1, NULL)"), "1")
         self.assertEqual(self._run("concat(true, false)"), "truefalse")

diff --git a/hogvm/typescript/src/__tests__/bytecode.test.ts b/hogvm/typescript/src/__tests__/bytecode.test.ts
@@ -2,7 +2,7 @@ import { executeHogQLBytecode, Operation as op } from '../bytecode'
 
 describe('HogQL Bytecode', () => {
     test('execution results', async () => {
-        const fields = { properties: { foo: 'bar' } }
+        const fields = { properties: { foo: 'bar', nullValue: null } }
         expect(await executeHogQLBytecode(['_h', op.INTEGER, 2, op.INTEGER, 1, op.PLUS], fields)).toBe(3)
         expect(await executeHogQLBytecode(['_h', op.INTEGER, 2, op.INTEGER, 1, op.MINUS], fields)).toBe(-1)
         expect(await executeHogQLBytecode(['_h', op.INTEGER, 2, op.INTEGER, 3, op.MULTIPLY], fields)).toBe(6)
@@ -62,6 +62,18 @@ describe('HogQL Bytecode', () => {
         expect(await executeHogQLBytecode(['_h', op.STRING, 'foo', op.STRING, 'properties', op.FIELD, 2], fields)).toBe(
             'bar'
         )
+        expect(
+            await executeHogQLBytecode(
+                ['_h', op.FALSE, op.STRING, 'foo', op.STRING, 'properties', op.FIELD, 2, op.CALL, 'ifNull', 2],
+                fields
+            )
+        ).toBe('bar')
+        expect(
+            await executeHogQLBytecode(
+                ['_h', op.FALSE, op.STRING, 'nullValue', op.STRING, 'properties', op.FIELD, 2, op.CALL, 'ifNull', 2],
+                fields
+            )
+        ).toBe(false)
         expect(
             await executeHogQLBytecode(['_h', op.STRING, 'another', op.STRING, 'arg', op.CALL, 'concat', 2], fields)
         ).toBe('arganother')

diff --git a/hogvm/typescript/src/bytecode.ts b/hogvm/typescript/src/bytecode.ts
@@ -213,6 +213,12 @@ export function executeHogQLBytecode(bytecode: any[], fields: Record<string, any
                 } else if (name == 'toFloat') {
                     const value = parseFloat(args[0])
                     stack.push(isNaN(value) ? null : value)
+                } else if (name == 'ifNull') {
+                    if (args[0] != null) {
+                        stack.push(args[0])
+                    } else {
+                        stack.push(args[1])
+                    }
                 } else {
                     throw new Error(`Unsupported function call: ${name}`)
                 }

diff --git a/posthog/hogql/test/test_bytecode.py b/posthog/hogql/test/test_bytecode.py
@@ -60,6 +60,10 @@ def test_bytecode_create(self):
             to_bytecode("concat('arg', 'another')"),
             [_H, op.STRING, "another", op.STRING, "arg", op.CALL, "concat", 2],
         )
+        self.assertEqual(
+            to_bytecode("ifNull(properties.email, false)"),
+            [_H, op.FALSE, op.STRING, "email", op.STRING, "properties", op.FIELD, 2, op.CALL, "ifNull", 2],
+        )
         self.assertEqual(to_bytecode("1 = 2"), [_H, op.INTEGER, 2, op.INTEGER, 1, op.EQ])
         self.assertEqual(to_bytecode("1 == 2"), [_H, op.INTEGER, 2, op.INTEGER, 1, op.EQ])
         self.assertEqual(to_bytecode("1 != 2"), [_H, op.INTEGER, 2, op.INTEGER, 1, op.NOT_EQ])