Fix crash in Bridge.exportToArrow when vector is a MapVector wrapped …

…in a ConstantVector and ArrowOptions.flattenConstant is true
facebookincubator · Dec 22, 2024 · d052924 · d052924
1 parent e9bb6c1
commit d052924
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/velox/vector/arrow/Bridge.cpp b/velox/vector/arrow/Bridge.cpp
@@ -1436,6 +1436,9 @@ void exportToArrow(
         0, newArrowSchema("i", "run_ends"), arrowSchema);
     bridgeHolder->setChildAtIndex(1, std::move(valuesChild), arrowSchema);
   } else {
+    if (vec->encoding() == VectorEncoding::Simple::CONSTANT && options.flattenConstant) {
+      VELOX_CHECK(vec->isScalar(), "Flattening is only supported for scalar types.");
+    }
     arrowSchema.format =
         exportArrowFormatStr(type, options, bridgeHolder->formatBuffer);
     arrowSchema.dictionary = nullptr;

diff --git a/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp b/velox/vector/arrow/tests/ArrowBridgeSchemaTest.cpp
@@ -252,6 +252,7 @@ TEST_F(ArrowBridgeSchemaExportTest, nested) {
 
   // Map.
   testNestedType(MAP(INTEGER(), DOUBLE()));
+  testNestedType(MAP(INTEGER(), DOUBLE()));
   testNestedType(MAP(VARBINARY(), BOOLEAN()));
   testNestedType(MAP(VARBINARY(), MAP(SMALLINT(), REAL())));
   testNestedType(MAP(VARBINARY(), MAP(SMALLINT(), MAP(INTEGER(), BIGINT()))));
@@ -290,6 +291,9 @@ TEST_F(ArrowBridgeSchemaExportTest, constant) {
 
   testConstant(ARRAY(INTEGER()), "+l");
   testConstant(ARRAY(UNKNOWN()), "+l");
+  VELOX_ASSERT_THROW(
+      testConstant(MAP(BOOLEAN(), REAL()), "+m", {false, true}),
+      "Flattening is only supported for scalar types.");
   testConstant(MAP(BOOLEAN(), REAL()), "+m");
   testConstant(MAP(UNKNOWN(), REAL()), "+m");
   testConstant(ROW({TIMESTAMP(), DOUBLE()}), "+s");