[arrow-flight-sql-postgresql] branch main updated: Add support for Int8 (#64)

[kou]

This is an automated email from the ASF dual-hosted git repository.

kou pushed a commit to branch main
in repository 
https://gitbox.apache.org/repos/asf/arrow-flight-sql-postgresql.git


The following commit(s) were added to refs/heads/main by this push:
     new 79d536e  Add support for Int8 (#64)
79d536e is described below

commit 79d536eb55a1a0b4fae62d41e0a24af28455c262
Author: Sutou Kouhei <k...@clear-code.com>
AuthorDate: Tue Aug 22 14:20:16 2023 +0900

    Add support for Int8 (#64)
    
    Closes GH-47
---
 src/afs.cc              |  9 +++++++++
 test/test-flight-sql.rb | 24 ++++++++++++++----------
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/afs.cc b/src/afs.cc
index c3befe9..9082336 100644
--- a/src/afs.cc
+++ b/src/afs.cc
@@ -732,6 +732,12 @@ class ArrowPGValueConverter : public arrow::ArrayVisitor {
        {
        }
 
+       arrow::Status Visit(const arrow::Int8Array& array)
+       {
+               datum_ = Int8GetDatum(array.Value(i_row_));
+               return arrow::Status::OK();
+       }
+
        arrow::Status Visit(const arrow::Int32Array& array)
        {
                datum_ = Int32GetDatum(array.Value(i_row_));
@@ -815,6 +821,9 @@ class PreparedStatement {
                {
                        switch (field->type()->id())
                        {
+                               case arrow::Type::INT8:
+                                       pgTypes.push_back(INT2OID);
+                                       break;
                                case arrow::Type::INT32:
                                        pgTypes.push_back(INT4OID);
                                        break;
diff --git a/test/test-flight-sql.rb b/test/test-flight-sql.rb
index cb995b4..c4ee2e5 100644
--- a/test/test-flight-sql.rb
+++ b/test/test-flight-sql.rb
@@ -75,30 +75,34 @@ SELECT * FROM data
     RESULT
   end
 
-  def test_insert_int32
+  data("int8", ["smallint", Arrow::Int8Array, [1, -2, 3]])
+  data("int32", ["integer", Arrow::Int32Array, [1, -2, 3]])
+  def test_insert_type
     unless flight_sql_client.respond_to?(:prepare)
       omit("red-arrow-flight-sql 14.0.0 or later is required")
     end
 
-    run_sql("CREATE TABLE data (value integer)")
+    pg_type, array_class, values = data
+    run_sql("CREATE TABLE data (value #{pg_type})")
 
     flight_sql_client.prepare("INSERT INTO data VALUES ($1)",
                               @options) do |statement|
-      values = Arrow::Int32Array.new([1, -2, 3])
-      statement.record_batch = Arrow::RecordBatch.new(value: values)
+      array = array_class.new(values)
+      statement.record_batch = Arrow::RecordBatch.new(value: array)
       n_changed_records = statement.execute_update(@options)
       assert_equal(3, n_changed_records)
     end
 
-    assert_equal([<<-RESULT, ""], run_sql("SELECT * FROM data"))
+    output = <<-RESULT
 SELECT * FROM data
  value 
 -------
-     1
-    -2
-     3
-(3 rows)
-
     RESULT
+    values.each do |value|
+      output << (" %5d\n" % value)
+    end
+    output << "(#{values.size} rows)\n"
+    output << "\n"
+    assert_equal([output, ""], run_sql("SELECT * FROM data"))
   end
 end