Implement OID reader (#235)

innobi · Jan 21, 2024 · cc9fe23 · cc9fe23
1 parent 3bd30b1
commit cc9fe23
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 3 deletions.
diff --git a/pantab/src/pantab.cpp b/pantab/src/pantab.cpp
@@ -464,7 +464,7 @@ class ReadHelper {
   struct ArrowArray *array_;
 };
 
-class IntegralReadHelper : public ReadHelper {
+template <typename T> class IntegralReadHelper : public ReadHelper {
   using ReadHelper::ReadHelper;
 
   auto Read(const hyperapi::Value &value) -> void override {
@@ -474,7 +474,7 @@ class IntegralReadHelper : public ReadHelper {
       }
       return;
     }
-    if (ArrowArrayAppendInt(array_, value.get<int64_t>())) {
+    if (ArrowArrayAppendInt(array_, value.get<T>())) {
       throw std::runtime_error("ArrowAppendInt failed");
     };
   }
@@ -629,9 +629,13 @@ static auto makeReadHelper(const ArrowSchemaView *schema_view,
     -> std::unique_ptr<ReadHelper> {
   switch (schema_view->type) {
   case NANOARROW_TYPE_INT16:
+    return std::unique_ptr<ReadHelper>(new IntegralReadHelper<int16_t>(array));
   case NANOARROW_TYPE_INT32:
+    return std::unique_ptr<ReadHelper>(new IntegralReadHelper<int32_t>(array));
   case NANOARROW_TYPE_INT64:
-    return std::unique_ptr<ReadHelper>(new IntegralReadHelper(array));
+    return std::unique_ptr<ReadHelper>(new IntegralReadHelper<int64_t>(array));
+  case NANOARROW_TYPE_UINT32:
+    return std::unique_ptr<ReadHelper>(new IntegralReadHelper<uint32_t>(array));
   case NANOARROW_TYPE_DOUBLE:
     return std::unique_ptr<ReadHelper>(new FloatReadHelper(array));
   case NANOARROW_TYPE_LARGE_BINARY:
@@ -659,6 +663,7 @@ static auto arrowTypeFromHyper(const hyperapi::SqlType &sqltype)
         case hyperapi::TypeTag::SmallInt : return NANOARROW_TYPE_INT16;
         case hyperapi::TypeTag::Int : return NANOARROW_TYPE_INT32;
         case hyperapi::TypeTag::BigInt : return NANOARROW_TYPE_INT64;
+        case hyperapi::TypeTag::Oid : return NANOARROW_TYPE_UINT32;
         case hyperapi::TypeTag::Double : return NANOARROW_TYPE_DOUBLE;
         case hyperapi::TypeTag::Bytes : return NANOARROW_TYPE_LARGE_BINARY;
         case hyperapi::TypeTag::Varchar : case hyperapi::TypeTag::Char :

diff --git a/pantab/tests/test_reader.py b/pantab/tests/test_reader.py
@@ -122,3 +122,39 @@ def test_read_varchar(tmp_hyper):
 
     result = pantab.frame_from_hyper(tmp_hyper, table=table_name)
     tm.assert_frame_equal(result, expected)
+
+
+def test_read_oid(tmp_hyper):
+    column_name = "OID Column"
+    table_name = tab_api.TableName("public", "table")
+    table = tab_api.TableDefinition(
+        table_name=table_name,
+        columns=[
+            tab_api.TableDefinition.Column(
+                name=column_name,
+                type=tab_api.SqlType.oid(),
+                nullability=tab_api.NOT_NULLABLE,
+            )
+        ],
+    )
+
+    with tab_api.HyperProcess(
+        telemetry=tab_api.Telemetry.DO_NOT_SEND_USAGE_DATA_TO_TABLEAU
+    ) as hyper:
+        with tab_api.Connection(
+            endpoint=hyper.endpoint,
+            database=tmp_hyper,
+            create_mode=tab_api.CreateMode.CREATE_AND_REPLACE,
+        ) as connection:
+            connection.catalog.create_table(table_definition=table)
+
+            with tab_api.Inserter(connection, table) as inserter:
+                inserter.add_rows([[123], [456]])
+                inserter.execute()
+
+    expected = pd.DataFrame(
+        [[123], [456]], columns=[column_name], dtype="uint32[pyarrow]"
+    )
+
+    result = pantab.frame_from_hyper(tmp_hyper, table=table_name)
+    tm.assert_frame_equal(result, expected)