diff --git a/velox/expression/ComplexWriterTypes.h b/velox/expression/ComplexWriterTypes.h index 4a2539f19a3d..c7cb91c233ae 100644 --- a/velox/expression/ComplexWriterTypes.h +++ b/velox/expression/ComplexWriterTypes.h @@ -765,7 +765,7 @@ class MapWriter { keysVector_ = &keysWriter_->vector(); valuesVector_ = &valuesWriter_->vector(); - innerOffset_ = std::max(keysVector_->size(), valuesVector_->size()); + innerOffset_ = std::min(keysVector_->size(), valuesVector_->size()); // Keys can never be null. keysVector_->resetNulls(); @@ -773,10 +773,7 @@ class MapWriter { keysWriter_->ensureSize(1); valuesWriter_->ensureSize(1); - VELOX_DCHECK( - keysVector_->size() == valuesVector_->size(), - "expect map keys and value vector sized to be synchronized"); - capacity_ = keysVector_->size(); + capacity_ = std::min(keysVector_->size(), valuesVector_->size()); } key_element_t& lastKeyWriter() { diff --git a/velox/expression/tests/MapWriterTest.cpp b/velox/expression/tests/MapWriterTest.cpp index cee1b59c933e..7eec5053b0fb 100644 --- a/velox/expression/tests/MapWriterTest.cpp +++ b/velox/expression/tests/MapWriterTest.cpp @@ -645,6 +645,64 @@ TEST_F(MapWriterTest, appendToKeysAndValues) { ASSERT_EQ(values->asFlatVector()->valueAt(1), 20); } +// Make sure MapWriter correctly handles 'result' MapVector, where keys and +// values child vectors have different sizes. +TEST_F(MapWriterTest, differentKeyValueSizes) { + using out_t = Map; + auto mapType = CppToType::create(); + + // Keys vector is shorter than values vector. + auto result = std::make_shared( + pool(), + mapType, + nullptr, // no nulls + 3, + makeIndices({0, 2, 4}), // offsets + makeIndices({2, 2, 2}), // sizes + makeFlatVector({0, 1, 0, 1, 0, 1}), // keys + makeFlatVector({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}) // values + ); + + EXPECT_EQ(6, result->mapKeys()->size()); + EXPECT_EQ(10, result->mapValues()->size()); + + { + // Write map at offset 0. + exec::VectorWriter vectorWriter; + vectorWriter.init(*result); + vectorWriter.setOffset(0); + auto& mapWriter = vectorWriter.current(); + mapWriter.copy_from(folly::F14FastMap{{1, 2}}); + + vectorWriter.commit(); + vectorWriter.finish(); + } + + { + // Write map at offset 2. + exec::VectorWriter vectorWriter; + vectorWriter.init(*result); + vectorWriter.setOffset(2); + auto& mapWriter = vectorWriter.current(); + mapWriter.copy_from(folly::F14FastMap{{3, 4}}); + + vectorWriter.commit(); + vectorWriter.finish(); + } + + // Verify sizes of keys and values vectors. + EXPECT_EQ(8, result->mapKeys()->size()); + EXPECT_EQ(8, result->mapValues()->size()); + + auto expected = makeMapVector({ + {{1, 2}}, + {{0, 2}, {1, 3}}, + {{3, 4}}, + }); + + assertEqualVectors(expected, result); +} + // Make sure copy from MapView correctly resizes children vectors. TEST_F(MapWriterTest, copyFromViewTypeResizedChildren) { using out_t = Map>;