Properly offset varlen outputs for multiple projected outputs in one row

yma11 · Jul 6, 2021 · 6d37f5d · 6d37f5d
1 parent 7c03dbf
commit 6d37f5d
Show file tree

Hide file tree

Showing 5 changed files with 94 additions and 23 deletions.
diff --git a/QueryEngine/Descriptors/QueryMemoryDescriptor.cpp b/QueryEngine/Descriptors/QueryMemoryDescriptor.cpp
@@ -1271,3 +1271,20 @@ std::optional<size_t> QueryMemoryDescriptor::varlenOutputBufferElemSize() const
   }
   return buffer_element_size;
 }
+
+size_t QueryMemoryDescriptor::varlenOutputRowSizeToSlot(const size_t slot_idx) const {
+  int64_t buffer_element_size{0};
+  CHECK_LT(slot_idx, col_slot_context_.getSlotCount());
+  for (size_t i = 0; i < slot_idx; i++) {
+    try {
+      const auto slot_element_size = col_slot_context_.varlenOutputElementSize(i);
+      if (slot_element_size < 0) {
+        continue;
+      }
+      buffer_element_size += slot_element_size;
+    } catch (...) {
+      continue;
+    }
+  }
+  return buffer_element_size;
+}
diff --git a/QueryEngine/Descriptors/QueryMemoryDescriptor.h b/QueryEngine/Descriptors/QueryMemoryDescriptor.h
@@ -333,6 +333,10 @@ class QueryMemoryDescriptor {
   // bump allocator
   std::optional<size_t> varlenOutputBufferElemSize() const;
 
+  // returns the number of bytes needed for all slots preceeding slot_idx. Used to compute
+  // the offset into the varlen buffer for each projected target in a given row.
+  size_t varlenOutputRowSizeToSlot(const size_t slot_idx) const;
+
   bool slotIsVarlenOutput(const size_t slot_idx) const {
     return col_slot_context_.slotIsVarlen(slot_idx);
   }

diff --git a/QueryEngine/RuntimeFunctions.cpp b/QueryEngine/RuntimeFunctions.cpp
@@ -380,6 +380,7 @@ extern "C" ALWAYS_INLINE int8_t* agg_id_varlen(int8_t* varlen_buffer,
                                                const int64_t offset,
                                                const int8_t* value,
                                                const int64_t size_bytes) {
+  printf("Offset: %ld\n", offset);
   for (auto i = 0; i < size_bytes; i++) {
     varlen_buffer[offset + i] = value[i];
   }

diff --git a/QueryEngine/TargetExprBuilder.cpp b/QueryEngine/TargetExprBuilder.cpp
@@ -382,8 +382,15 @@ void TargetExprCodegen::codegenAggregate(
       const auto output_buffer_slot = LL_BUILDER.CreateZExt(
           LL_BUILDER.CreateLoad(get_arg_by_name(ROW_FUNC, "old_total_matched")),
           llvm::Type::getInt64Ty(LL_CONTEXT));
-      const auto output_buffer_slot_bytes = LL_BUILDER.CreateMul(
-          output_buffer_slot, executor->cgen_state_->llInt(chosen_bytes));
+      const auto varlen_buffer_row_sz = query_mem_desc.varlenOutputBufferElemSize();
+      CHECK(varlen_buffer_row_sz);
+      const auto output_buffer_slot_bytes = LL_BUILDER.CreateAdd(
+          LL_BUILDER.CreateMul(
+              output_buffer_slot,
+              executor->cgen_state_->llInt(static_cast<int64_t>(*varlen_buffer_row_sz))),
+          executor->cgen_state_->llInt(static_cast<int64_t>(
+              query_mem_desc.varlenOutputRowSizeToSlot(slot_index))));
+      ;
       std::vector<llvm::Value*> varlen_agg_args{
           executor->castToIntPtrTyIn(varlen_output_buffer, 8),
           output_buffer_slot_bytes,

diff --git a/Tests/GeospatialTest.cpp b/Tests/GeospatialTest.cpp
@@ -603,27 +603,6 @@ TEST_P(GeoSpatialTestTablesFixture, Basics) {
                                    "ST_Distance(ST_GeomFromText('POINT(0 0)'), p) < 1;",
                                    dt,
                                    false))));
-    ASSERT_EQ(
-        "POINT (2 2)",
-        boost::get<std::string>(v<NullableString>(run_simple_agg(
-            "SELECT ST_GeomFromText('POINT(2 2)') FROM geospatial_test WHERE id = 2;",
-            dt,
-            false))));
-    ASSERT_EQ(
-        "POINT (2 2)",
-        boost::get<std::string>(v<NullableString>(run_simple_agg(
-            "SELECT ST_Point(2,2) FROM geospatial_test WHERE id = 2;", dt, false))));
-    // requires punt to CPU
-    SKIP_ON_AGGREGATOR(ASSERT_EQ(
-        "POINT (2 2)",
-        boost::get<std::string>(v<NullableString>(run_simple_agg(
-            "SELECT ST_Point(id,id) FROM geospatial_test WHERE id = 2;", dt, false)))));
-    SKIP_ON_AGGREGATOR(ASSERT_EQ(
-        "POINT (2 2)",
-        boost::get<std::string>(v<NullableString>(run_simple_agg(
-            "SELECT ST_SetSRID(ST_Point(id,id),4326) FROM geospatial_test WHERE id = 2;",
-            dt,
-            false)))));
 
     // more distance
     ASSERT_NEAR(
@@ -996,6 +975,56 @@ TEST_P(GeoSpatialTestTablesFixture, Basics) {
   }
 }
 
+TEST_P(GeoSpatialTestTablesFixture, Constructors) {
+  for (auto dt : {ExecutorDeviceType::CPU, ExecutorDeviceType::GPU}) {
+    SKIP_NO_GPU();
+
+    {
+      auto rows = run_multiple_agg(
+          R"(SELECT ST_Point(id, id), id, ST_Point(id + 1, id + 2) FROM geospatial_test WHERE id < 2 ORDER BY 2;)",
+          dt);
+      rows->setGeoReturnType(ResultSet::GeoReturnType::GeoTargetValue);
+      EXPECT_EQ(rows->rowCount(), size_t(2));
+
+      auto process_row = [&](const int64_t id_for_row) {
+        auto row = rows->getNextRow(false, false);
+        EXPECT_EQ(row.size(), size_t(3));
+        const auto id = v<int64_t>(row[1]);
+        EXPECT_EQ(id, id_for_row);
+        const auto first_geo_tv = boost::get<GeoTargetValue>(row[0]);
+        CHECK(first_geo_tv);
+        const auto first_pt = boost::get<GeoPointTargetValue>(*first_geo_tv);
+        CHECK(first_pt.coords->data());
+        const double* first_pt_array = first_pt.coords->data();
+        EXPECT_EQ(first_pt_array[0], static_cast<double>(id));
+        EXPECT_EQ(first_pt_array[1], static_cast<double>(id));
+
+        const auto second_geo_tv = boost::get<GeoTargetValue>(row[2]);
+        CHECK(second_geo_tv);
+        const auto second_pt = boost::get<GeoPointTargetValue>(*second_geo_tv);
+        CHECK(second_pt.coords->data());
+        const double* second_pt_array = second_pt.coords->data();
+        EXPECT_EQ(second_pt_array[0], static_cast<double>(id + 1));
+        EXPECT_EQ(second_pt_array[1], static_cast<double>(id + 2));
+      };
+
+      process_row(0);
+      process_row(1);
+    }
+
+    ASSERT_EQ(
+        "POINT (2 2)",
+        boost::get<std::string>(v<NullableString>(run_simple_agg(
+            "SELECT ST_Point(id,id) FROM geospatial_test WHERE id = 2;", dt, false))));
+    ASSERT_EQ(
+        "POINT (2 2)",
+        boost::get<std::string>(v<NullableString>(run_simple_agg(
+            "SELECT ST_SetSRID(ST_Point(id,id),4326) FROM geospatial_test WHERE id = 2;",
+            dt,
+            false))));
+  }
+}
+
 INSTANTIATE_TEST_SUITE_P(GeoSpatialTablesTests,
                          GeoSpatialTestTablesFixture,
                          ::testing::Values(true, false));
@@ -1793,6 +1822,19 @@ TEST(GeoSpatial, Math) {
   }
 }
 
+TEST(GeoSpatial, Projections) {
+  for (auto dt : {ExecutorDeviceType::CPU, ExecutorDeviceType::GPU}) {
+    SKIP_NO_GPU();
+
+    ASSERT_EQ("POINT (2 2)",
+              boost::get<std::string>(v<NullableString>(
+                  run_simple_agg("SELECT ST_GeomFromText('POINT(2 2)');", dt, false))));
+    ASSERT_EQ("POINT (2 2)",
+              boost::get<std::string>(
+                  v<NullableString>(run_simple_agg("SELECT ST_Point(2,2);", dt, false))));
+  }
+}
+
 class GeoSpatialTempTables : public ::testing::Test {
  protected:
   void SetUp() override { import_geospatial_test(/*with_temporary_tables=*/true); }