add support for operator.eq and operator.ne for TextEncodingNone

Author: guilhermeleobas

rbc/heavyai/buffer.py

@@ -51,6 +51,10 @@ class OmnisciBufferType(typesystem.Type):
     def pass_by_value(self):
         return False
 
+    @property
+    def numba_pointer_type(self):
+        return BufferPointer
+
     @classmethod
     def preprocess_args(cls, args):
         assert len(args) == 1, args
@@ -78,11 +82,11 @@ def tonumba(self, bool_is_int8=None):
             *extra_members
         )
         buffer_type._params['NumbaType'] = BufferType
-        buffer_type._params['NumbaPointerType'] = BufferPointer
+        buffer_type._params['NumbaPointerType'] = self.numba_pointer_type
         numba_type = buffer_type.tonumba(bool_is_int8=True)
         if self.pass_by_value:
             return numba_type
-        return BufferPointer(numba_type)
+        return self.numba_pointer_type(numba_type)
 
 
 class BufferType(types.Type):

rbc/heavyai/pipeline.py

@@ -158,7 +158,7 @@ def define_pipelines(self):
         # namely the "nopython" pipeline
         pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
         # Add the new pass to run after IRProcessing
-        # pm.add_pass_after(AutoFreeBuffers, IRProcessing)
+        pm.add_pass_after(AutoFreeBuffers, IRProcessing)
         pm.add_pass_after(CheckRaiseStmts, IRProcessing)
         pm.add_pass_after(DTypeComparison, ReconstructSSA)
         # finalize

rbc/heavyai/text_encoding_none.py

@@ -1,19 +1,28 @@
 '''Omnisci TextEncodingNone type that corresponds to Omnisci type TEXT ENCODED NONE.
 '''
 
-__all__ = ['TextEncodingNonePointer', 'TextEncodingNone', 'OmnisciTextEncodingNoneType']
+__all__ = ['TextEncodingNonePointer', 'TextEncodingNone', 'HeavyDBTextEncodingNoneType']
 
+import operator
 from rbc import typesystem
+from rbc.targetinfo import TargetInfo
+from rbc.errors import RequireLiteralValue
 from .buffer import (
     BufferPointer, Buffer, OmnisciBufferType,
     omnisci_buffer_constructor)
-from numba.core import types, extending
+from numba.core import types, extending, cgutils
+from llvmlite import ir
+from typing import Union
 
 
-class OmnisciTextEncodingNoneType(OmnisciBufferType):
+class HeavyDBTextEncodingNoneType(OmnisciBufferType):
     """Omnisci TextEncodingNone type for RBC typesystem.
     """
 
+    @property
+    def numba_pointer_type(self):
+        return TextEncodingNonePointer
+
     @classmethod
     def preprocess_args(cls, args):
         element_type = typesystem.Type.fromstring('char8')
@@ -32,7 +41,8 @@ def match(self, other):
             return 2
 
 
-TextEncodingNonePointer = BufferPointer
+class TextEncodingNonePointer(BufferPointer):
+    pass
 
 
 class TextEncodingNone(Buffer):
@@ -53,25 +63,91 @@ class TextEncodingNone(Buffer):
 
         from rbc.heavydb import TextEncodingNone
 
-        @omnisci('TextEncodingNone(int32, int32)')
+        @heavydb('TextEncodingNone(int32, int32)')
         def make_abc(first, n):
             r = TextEncodingNone(n)
             for i in range(n):
                 r[i] = first + i
             return r
+
+
+    .. code-block:: python
+
+        from rbc.heavydb import TextEncodingNone
+        @heavydb('TextEncodingNone()')
+        def make_text():
+            return TextEncodingNone('some text here')
+
     '''
 
-    def __init__(self, size: int):
+    def __init__(self, size: Union[int, str]):
         pass
 
 
+@extending.overload(operator.eq)
+def text_encoding_none_eq(a, b):
+    if isinstance(a, TextEncodingNonePointer) and isinstance(b, TextEncodingNonePointer):
+
+        def impl(a, b):
+            if len(a) != len(b):
+                return False
+            for i in range(0, len(a)):
+                if a[i] != b[i]:
+                    return False
+            return True
+        return impl
+    elif isinstance(a, TextEncodingNonePointer) and isinstance(b, types.StringLiteral):
+        lv = b.literal_value
+        sz = len(lv)
+
+        def impl(a, b):
+            if len(a) != sz:
+                return False
+            t = TextEncodingNone(lv)
+            return a == t
+        return impl
+
+
+@extending.overload(operator.ne)
+def text_encoding_none_ne(a, b):
+    if isinstance(a, TextEncodingNonePointer):
+        if isinstance(b, (TextEncodingNonePointer, types.StringLiteral)):
+            def impl(a, b):
+                return not(a == b)
+            return impl
+
+
 @extending.lower_builtin(TextEncodingNone, types.Integer)
 def omnisci_text_encoding_none_constructor(context, builder, sig, args):
     return omnisci_buffer_constructor(context, builder, sig, args)
 
 
+@extending.lower_builtin(TextEncodingNone, types.StringLiteral)
+def omnisci_text_encoding_none_constructor_literal(context, builder, sig, args):
+    int64_t = ir.IntType(64)
+    int8_t_ptr = ir.IntType(8).as_pointer()
+
+    literal_value = sig.args[0].literal_value
+    sz = int64_t(len(literal_value))
+
+    # arr = {ptr, size, is_null}*
+    arr = omnisci_buffer_constructor(context, builder, sig.return_type(types.int64), [sz])
+    ptr = builder.extract_value(builder.load(arr), [0])
+
+    msg_bytes = literal_value.encode('utf-8')
+    msg_const = cgutils.make_bytearray(msg_bytes)
+    msg_global_var = cgutils.global_constant(builder.module, f"Text({literal_value})", msg_const)
+    msg_ptr = builder.bitcast(msg_global_var, int8_t_ptr)
+    sizeof_char = TargetInfo().sizeof('char')
+    cgutils.raw_memcpy(builder, ptr, msg_ptr, sz, sizeof_char)
+    return arr
+
+
 @extending.type_callable(TextEncodingNone)
 def type_omnisci_text_encoding_none(context):
-    def typer(size):
-        return typesystem.Type.fromobject('TextEncodingNone').tonumba()
+    def typer(arg):
+        if isinstance(arg, types.UnicodeType):
+            raise RequireLiteralValue()
+        if isinstance(arg, (types.Integer, types.StringLiteral)):
+            return typesystem.Type.fromobject('TextEncodingNone').tonumba()
     return typer

rbc/heavydb.py

@@ -14,7 +14,7 @@
 from .thrift import Client as ThriftClient
 from . import heavyai
 from .heavyai import (
-    OmnisciArrayType, OmnisciTextEncodingNoneType, OmnisciTextEncodingDictType,
+    OmnisciArrayType, HeavyDBTextEncodingNoneType, OmnisciTextEncodingDictType,
     OmnisciOutputColumnType, OmnisciColumnType,
     OmnisciCompilerPipeline, OmnisciCursorType,
     BufferMeta, OmnisciColumnListType, OmnisciTableFunctionManagerType)
@@ -398,7 +398,7 @@ def add(a, b):
         Constant='int32|sizer=Constant',
         PreFlight='int32|sizer=PreFlight',
         ColumnList='OmnisciColumnListType',
-        TextEncodingNone='OmnisciTextEncodingNoneType',
+        TextEncodingNone='HeavyDBTextEncodingNoneType',
         TextEncodingDict='OmnisciTextEncodingDictType',
         TableFunctionManager='OmnisciTableFunctionManagerType<>',
         UDTF='int32|kind=UDTF'
@@ -894,7 +894,7 @@ def _get_ext_arguments_map(self):
             ext_arguments_map['OmnisciOutputColumnListType<%s>' % ptr_type] \
                 = ext_arguments_map.get('ColumnList<%s>' % T)
 
-        ext_arguments_map['OmnisciTextEncodingNoneType<char8>'] = \
+        ext_arguments_map['HeavyDBTextEncodingNoneType<char8>'] = \
             ext_arguments_map.get('TextEncodingNone')
 
         values = list(ext_arguments_map.values())
@@ -1398,7 +1398,7 @@ def format_type(self, typ: typesystem.Type):
         elif isinstance(typ, OmnisciCursorType):
             p = tuple(map(self.format_type, typ[0]))
             typ = typesystem.Type(('Cursor',) + p, **typ._params)
-        elif isinstance(typ, OmnisciTextEncodingNoneType):
+        elif isinstance(typ, HeavyDBTextEncodingNoneType):
             typ = typ.copy().params(typename='TextEncodingNone')
             use_typename = True
         elif isinstance(typ, OmnisciTextEncodingDictType):
@@ -1444,7 +1444,7 @@ def remote_call(self, func, ftype: typesystem.Type, arguments: tuple, hold=False
 
             if isinstance(atype, (OmnisciColumnType, OmnisciColumnListType)):
                 args.append(f'CURSOR({a})')
-            elif isinstance(atype, OmnisciTextEncodingNoneType):
+            elif isinstance(atype, HeavyDBTextEncodingNoneType):
                 if isinstance(a, bytes):
                     a = repr(a.decode())
                 elif isinstance(a, str):

rbc/tests/heavyai/test_text.py

@@ -135,3 +135,53 @@ def myupper(s):
 
     for n, u in result:
         assert n.upper() == u
+
+
+def test_TextEncodingNone_str_constructor(heavydb):
+    heavydb.reset()
+
+    @heavydb('TextEncodingNone(int32)')
+    def constructor(_):
+        return TextEncodingNone('hello world')
+
+    assert constructor(3).execute() == 'hello world'
+
+
+def test_TextEncodingNone_eq(heavydb):
+
+    heavydb.reset()
+
+    @heavydb('int32(TextEncodingNone, TextEncodingNone)')
+    def eq1(a, b):
+        return a == b
+
+    @heavydb('int32(TextEncodingNone)')
+    def eq2(a):
+        return a == 'world'
+
+    assert eq1('hello', 'hello').execute() == 1
+    assert eq1('c', 'c').execute() == 1
+    assert eq1('hello', 'h').execute() == 0
+    assert eq1('hello', 'hello2').execute() == 0
+
+    assert eq2('world').execute() == 1
+
+
+def test_TextEncodingNone_ne(heavydb):
+
+    heavydb.reset()
+
+    @heavydb('int32(TextEncodingNone, TextEncodingNone)')
+    def ne1(a, b):
+        return a != b
+
+    @heavydb('int32(TextEncodingNone)')
+    def ne2(a):
+        return a != 'world'
+
+    assert ne1('hello', 'hello').execute() == 0
+    assert ne1('c', 'c').execute() == 0
+    assert ne1('hello', 'h').execute() == 1
+    assert ne1('hello', 'hello2').execute() == 1
+
+    assert ne2('world').execute() == 0