Implement int/uint/float types, except for float16

Signed-off-by: Deepyaman Datta <deepyaman.datta@utexas.edu>
unionai-oss · Feb 9, 2025 · 5352be0 · 5352be0
1 parent 213ccae
commit 5352be0
Show file tree

Hide file tree

Showing 2 changed files with 121 additions and 0 deletions.
diff --git a/pandera/engines/ibis_engine.py b/pandera/engines/ibis_engine.py
@@ -83,6 +83,26 @@ def dtype(cls, data_type: Any) -> dtypes.DataType:
 ###############################################################################
 
 
+@Engine.register_dtype(
+    equivalents=[np.int8, dtypes.Int8, dtypes.Int8(), dt.Int8, dt.int8]
+)
+@immutable
+class Int8(DataType, dtypes.Int8):
+    """Semantic representation of a :class:`dt.Int8`."""
+
+    type = dt.int8
+
+
+@Engine.register_dtype(
+    equivalents=[np.int16, dtypes.Int16, dtypes.Int16(), dt.Int16, dt.int16]
+)
+@immutable
+class Int16(DataType, dtypes.Int16):
+    """Semantic representation of a :class:`dt.Int16`."""
+
+    type = dt.int16
+
+
 @Engine.register_dtype(
     equivalents=[np.int32, dtypes.Int32, dtypes.Int32(), dt.Int32, dt.int32]
 )
@@ -110,11 +130,88 @@ class Int64(DataType, dtypes.Int64):
     type = dt.int64
 
 
+###############################################################################
+# unsigned integer
+###############################################################################
+
+
+@Engine.register_dtype(
+    equivalents=[np.uint8, dtypes.UInt8, dtypes.UInt8(), dt.UInt8, dt.uint8]
+)
+@immutable
+class UInt8(DataType, dtypes.UInt8):
+    """Semantic representation of a :class:`dt.UInt8`."""
+
+    type = dt.uint8
+
+
+@Engine.register_dtype(
+    equivalents=[
+        np.uint16,
+        dtypes.UInt16,
+        dtypes.UInt16(),
+        dt.UInt16,
+        dt.uint16,
+    ]
+)
+@immutable
+class UInt16(DataType, dtypes.UInt16):
+    """Semantic representation of a :class:`dt.UInt16`."""
+
+    type = dt.uint16
+
+
+@Engine.register_dtype(
+    equivalents=[
+        np.uint32,
+        dtypes.UInt32,
+        dtypes.UInt32(),
+        dt.UInt32,
+        dt.uint32,
+    ]
+)
+@immutable
+class UInt32(DataType, dtypes.UInt32):
+    """Semantic representation of a :class:`dt.UInt32`."""
+
+    type = dt.uint32
+
+
+@Engine.register_dtype(
+    equivalents=[
+        np.uint64,
+        dtypes.UInt64,
+        dtypes.UInt64(),
+        dt.UInt64,
+        dt.uint64,
+    ]
+)
+@immutable
+class UInt64(DataType, dtypes.UInt64):
+    """Semantic representation of a :class:`dt.UInt64`."""
+
+    type = dt.uint64
+
+
 ###############################################################################
 # float
 ###############################################################################
 
 
+@Engine.register_dtype(
+    equivalents=[
+        np.float32,
+        dtypes.Float32,
+        dtypes.Float32(),
+        dt.Float32,
+        dt.float32,
+    ]
+)
+@immutable
+class Float32(DataType, dtypes.Float32):
+    """Semantic representation of a :class:`dt.Float32`."""
+
+
 @Engine.register_dtype(
     equivalents=[
         float,

diff --git a/tests/ibis/test_ibis_builtin_checks.py b/tests/ibis/test_ibis_builtin_checks.py
@@ -149,9 +149,21 @@ def get_data_param(self):
         data types would be tested"""
         return {
             "test_equal_to_check": [
+                {"datatype": dt.UInt8, "data": self.sample_numeric_data},
+                {"datatype": dt.UInt16, "data": self.sample_numeric_data},
+                {"datatype": dt.UInt32, "data": self.sample_numeric_data},
+                {"datatype": dt.UInt64, "data": self.sample_numeric_data},
+                {"datatype": dt.Int8, "data": self.sample_numeric_data},
+                {"datatype": dt.Int16, "data": self.sample_numeric_data},
                 {"datatype": dt.Int32, "data": self.sample_numeric_data},
                 {"datatype": dt.Int64, "data": self.sample_numeric_data},
                 {"datatype": dt.String, "data": self.sample_string_data},
+                {
+                    "datatype": dt.Float32,
+                    "data": self.convert_data(
+                        self.sample_numeric_data, "float32"
+                    ),
+                },
                 {
                     "datatype": dt.Float64,
                     "data": self.convert_data(
@@ -206,9 +218,21 @@ def get_data_param(self):
         data types would be tested"""
         return {
             "test_not_equal_to_check": [
+                {"datatype": dt.UInt8, "data": self.sample_numeric_data},
+                {"datatype": dt.UInt16, "data": self.sample_numeric_data},
+                {"datatype": dt.UInt32, "data": self.sample_numeric_data},
+                {"datatype": dt.UInt64, "data": self.sample_numeric_data},
+                {"datatype": dt.Int8, "data": self.sample_numeric_data},
+                {"datatype": dt.Int16, "data": self.sample_numeric_data},
                 {"datatype": dt.Int32, "data": self.sample_numeric_data},
                 {"datatype": dt.Int64, "data": self.sample_numeric_data},
                 {"datatype": dt.String, "data": self.sample_string_data},
+                {
+                    "datatype": dt.Float32,
+                    "data": self.convert_data(
+                        self.sample_numeric_data, "float32"
+                    ),
+                },
                 {
                     "datatype": dt.Float64,
                     "data": self.convert_data(