Fix out of range access with 1M over records response (#155)

Closes GH-154

Author: kou

src/afs.cc

@@ -1032,40 +1032,42 @@ class ArrowArrayBuilder<ArrowType, arrow::enable_if_has_c_type<ArrowType>>
 
 	arrow::Status build_not_null(uint64_t iTuple, uint64_t iTupleEnd)
 	{
-		values_.resize(iTupleEnd - iTuple);
-		for (; iTuple < iTupleEnd; iTuple++)
+		uint64_t n = iTupleEnd - iTuple;
+		values_.resize(n);
+		for (uint64_t i = 0; i < n; ++i)
 		{
 			bool isNull;
-			auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple],
+			auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple + i],
 			                           SPI_tuptable->tupdesc,
 			                           iAttribute_ + 1,
 			                           &isNull);
-			values_[iTuple] = convert_value<ArrowType>(datum);
+			values_[i] = convert_value<ArrowType>(datum);
 		}
 		return builder_->AppendValues(values_.data(), values_.size());
 	}
 
 	arrow::Status build_may_null(uint64_t iTuple, uint64_t iTupleEnd)
 	{
 		bool haveNull = false;
-		values_.resize(iTupleEnd - iTuple);
-		validBytes_.resize(iTupleEnd - iTuple);
-		for (; iTuple < iTupleEnd; iTuple++)
+		uint64_t n = iTupleEnd - iTuple;
+		values_.resize(n);
+		validBytes_.resize(n);
+		for (uint64_t i = 0; i < n; ++i)
 		{
 			bool isNull;
-			auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple],
+			auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple + i],
 			                           SPI_tuptable->tupdesc,
 			                           iAttribute_ + 1,
 			                           &isNull);
 			if (isNull)
 			{
 				haveNull = true;
-				validBytes_[iTuple] = 0;
+				validBytes_[i] = 0;
 			}
 			else
 			{
-				validBytes_[iTuple] = 1;
-				values_[iTuple] = convert_value<ArrowType>(datum);
+				validBytes_[i] = 1;
+				values_[i] = convert_value<ArrowType>(datum);
 			}
 		}
 		if (haveNull)
@@ -1148,41 +1150,42 @@ class ArrowArrayBuilder<ArrowType, arrow::enable_if_base_binary<ArrowType>>
 
 	arrow::Status build_not_null(uint64_t iTuple, uint64_t iTupleEnd)
 	{
-		values_.resize(iTupleEnd - iTuple);
-		for (; iTuple < iTupleEnd; iTuple++)
+		uint64_t n = iTupleEnd - iTuple;
+		values_.resize(n);
+		for (uint64_t i = 0; i < n; ++i)
 		{
 			bool isNull;
-			auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple],
+			auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple + i],
 			                           SPI_tuptable->tupdesc,
 			                           iAttribute_ + 1,
 			                           &isNull);
-			values_[iTuple] = std::string(VARDATA_ANY(datum), VARSIZE_ANY_EXHDR(datum));
+			values_[i] = std::string(VARDATA_ANY(datum), VARSIZE_ANY_EXHDR(datum));
 		}
 		return builder_->AppendValues(values_);
 	}
 
 	arrow::Status build_may_null(uint64_t iTuple, uint64_t iTupleEnd)
 	{
 		bool haveNull = false;
-		values_.resize(iTupleEnd - iTuple);
-		validBytes_.resize(iTupleEnd - iTuple);
-		for (; iTuple < iTupleEnd; iTuple++)
+		uint64_t n = iTupleEnd - iTuple;
+		values_.resize(n);
+		validBytes_.resize(n);
+		for (uint64_t i = 0; i < n; ++i)
 		{
 			bool isNull;
-			auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple],
+			auto datum = SPI_getbinval(SPI_tuptable->vals[iTuple + i],
 			                           SPI_tuptable->tupdesc,
 			                           iAttribute_ + 1,
 			                           &isNull);
 			if (isNull)
 			{
 				haveNull = true;
-				validBytes_[iTuple] = 0;
+				validBytes_[i] = 0;
 			}
 			else
 			{
-				validBytes_[iTuple] = 1;
-				values_[iTuple] =
-					std::string(VARDATA_ANY(datum), VARSIZE_ANY_EXHDR(datum));
+				validBytes_[i] = 1;
+				values_[i] = std::string(VARDATA_ANY(datum), VARSIZE_ANY_EXHDR(datum));
 			}
 		}
 		if (haveNull)

test/helper/sandbox.rb

@@ -123,6 +123,7 @@ def running?
     end
 
     def initdb(shared_preload_libraries: [],
+               max_n_rows_per_record_batch: nil,
                db_path: "db",
                port: 25432,
                flight_sql_port: 35432)
@@ -163,6 +164,10 @@ def initdb(shared_preload_libraries: [],
         conf.puts("shared_preload_libraries = " +
                   "'#{shared_preload_libraries.join(",")}'")
         conf.puts("arrow_flight_sql.uri = '#{@flight_sql_uri}'")
+        if max_n_rows_per_record_batch
+          conf.puts("arrow_flight_sql.max_n_rows_per_record_batch = " +
+                    "#{max_n_rows_per_record_batch}")
+        end
         yield(conf) if block_given?
       end
       pg_hba_conf = File.join(@dir, "pg_hba.conf")
@@ -343,6 +348,7 @@ def setup_db
       @postgresql = PostgreSQL.new(@tmp_dir)
       options = {
         shared_preload_libraries: shared_preload_libraries,
+        max_n_rows_per_record_batch: max_n_rows_per_record_batch,
       }
       @postgresql.initdb(**options)
       yield
@@ -352,6 +358,10 @@ def shared_preload_libraries
       ["arrow_flight_sql"]
     end
 
+    def max_n_rows_per_record_batch
+      nil
+    end
+
     def start_postgres
       @postgresql.start
     end

test/test-flight-sql.rb

@@ -18,6 +18,10 @@
 class FlightSQLTest < Test::Unit::TestCase
   include Helper::Sandbox
 
+  def max_n_rows_per_record_batch
+    10
+  end
+
   setup do
     unless flight_client.respond_to?(:authenticate_basic)
       omit("ArrowFlight::Client#authenticate_basic is needed" )
@@ -123,6 +127,41 @@ def test_select_from
                  reader.read_all)
   end
 
+  data(:type, ["integer", "text"])
+  data(:may_null, [true, false])
+  def test_select_multiple_record_batches_response
+    type = data[:type]
+    may_null = data[:may_null]
+    if type == "integer"
+      arrow_type = Arrow::Int32DataType.new
+      values = 30.times.to_a
+    else
+      arrow_type = Arrow::StringDataType.new
+      values = 30.times.collect(&:to_s)
+    end
+
+    create_table = "CREATE TABLE data (value #{type}"
+    # XXX: "NOT NULL" doesn't change FormData_pg_attribute::attnotnull
+    # returned by "SELECT *". So the "may_null" data driven test
+    # parameter is meaningless...
+    create_table << " NOT NULL" unless may_null
+    create_table << ")"
+    run_sql(create_table)
+    run_sql("INSERT INTO data VALUES " +
+            values.collect {|value| "(#{value})"}.join(", "))
+
+    info = flight_sql_client.execute("SELECT * FROM data", @options)
+    assert_equal(Arrow::Schema.new(value: arrow_type),
+                 info.get_schema)
+    endpoint = info.endpoints.first
+    reader = flight_sql_client.do_get(endpoint.ticket, @options)
+    chunks = values.each_slice(max_n_rows_per_record_batch).collect do |chunk|
+      arrow_type.array_class.new(chunk)
+    end
+    assert_equal(Arrow::Table.new(value: Arrow::ChunkedArray.new(chunks)),
+                 reader.read_all)
+  end
+
   def test_insert_direct
     unless flight_sql_client.respond_to?(:execute_update)
       omit("red-arrow-flight-sql 13.0.0 or later is required")