Fix Array/Map/RowVector::copyRanges for UNKNOWN source (facebookincub…

…ator#6607)

Summary:
Fix crashes when copying from TypeKind::UNKNOWN vector in the following methods:
- Map/Array/RowVector::copyRanges
- FlatVector<StringView>::copyRanges
- FlatVector<bool>::copy(source, rows, toSourceRow)
- RowVector::copy(source, rows, toSourceRow)

Fixes facebookincubator#6606

Pull Request resolved: facebookincubator#6607

Reviewed By: xiaoxmeng

Differential Revision: D49348224

Pulled By: mbasmanova

fbshipit-source-id: 57598428a36787ee3c57a1b7c9f6cffd11206bf6

velox/vector/ComplexVector.cpp

@@ -180,6 +180,11 @@ void RowVector::copy(
     const BaseVector* source,
     const SelectivityVector& rows,
     const vector_size_t* toSourceRow) {
+  if (source->typeKind() == TypeKind::UNKNOWN) {
+    rows.applyToSelected([&](auto row) { setNull(row, true); });
+    return;
+  }
+
   // Copy non-null values.
   SelectivityVector nonNullRows = rows;
 
@@ -261,13 +266,47 @@ void RowVector::copy(
   }
 }
 
+namespace {
+
+// Runs quick checks to determine whether input vector has only null values.
+// @return true if vector has only null values; false if vector may have
+// non-null values.
+bool isAllNullVector(const BaseVector& vector) {
+  if (vector.typeKind() == TypeKind::UNKNOWN) {
+    return true;
+  }
+
+  if (vector.isConstantEncoding()) {
+    return vector.isNullAt(0);
+  }
+
+  auto leafVector = vector.wrappedVector();
+  if (leafVector->isConstantEncoding()) {
+    // A null constant does not have a value vector, so wrappedVector
+    // returns the constant.
+    VELOX_CHECK(leafVector->isNullAt(0));
+    return true;
+  }
+  return false;
+}
+} // namespace
+
 void RowVector::copyRanges(
     const BaseVector* source,
     const folly::Range<const CopyRange*>& ranges) {
   if (ranges.empty()) {
     return;
   }
 
+  if (isAllNullVector(*source)) {
+    for (const auto& range : ranges) {
+      for (auto i = 0; i < range.count; ++i) {
+        setNull(range.targetIndex + i, true);
+      }
+    }
+    return;
+  }
+
   auto minTargetIndex = std::numeric_limits<vector_size_t>::max();
   auto maxTargetIndex = std::numeric_limits<vector_size_t>::min();
   for (auto& r : ranges) {
@@ -422,23 +461,29 @@ void ArrayVectorBase::copyRangesImpl(
     const BaseVector* source,
     const folly::Range<const BaseVector::CopyRange*>& ranges,
     VectorPtr* targetValues,
-    const BaseVector* sourceValues,
-    VectorPtr* targetKeys,
-    const BaseVector* sourceKeys) {
-  auto sourceValue = source->wrappedVector();
-  if (sourceValue->isConstantEncoding()) {
-    // A null constant does not have a value vector, so wrappedVector
-    // returns the constant.
-    VELOX_CHECK(sourceValue->isNullAt(0));
-    for (auto& r : ranges) {
-      for (auto i = 0; i < r.count; ++i) {
-        setNull(r.targetIndex + i, true);
+    VectorPtr* targetKeys) {
+  if (isAllNullVector(*source)) {
+    for (const auto& range : ranges) {
+      for (auto i = 0; i < range.count; ++i) {
+        setNull(range.targetIndex + i, true);
       }
     }
     return;
   }
-  VELOX_CHECK_EQ(sourceValue->encoding(), encoding());
-  auto sourceArray = sourceValue->asUnchecked<ArrayVectorBase>();
+
+  const BaseVector* sourceValues;
+  const BaseVector* sourceKeys = nullptr;
+
+  auto leafSource = source->wrappedVector();
+  VELOX_CHECK_EQ(leafSource->encoding(), encoding());
+
+  if (typeKind_ == TypeKind::ARRAY) {
+    sourceValues = leafSource->as<ArrayVector>()->elements().get();
+  } else {
+    sourceValues = leafSource->as<MapVector>()->mapValues().get();
+    sourceKeys = leafSource->as<MapVector>()->mapKeys().get();
+  }
+
   if (targetKeys) {
     BaseVector::ensureWritable(
         SelectivityVector::empty(),
@@ -452,6 +497,8 @@ void ArrayVectorBase::copyRangesImpl(
         pool(),
         *targetValues);
   }
+
+  auto sourceArray = leafSource->asUnchecked<ArrayVectorBase>();
   auto setNotNulls = mayHaveNulls() || source->mayHaveNulls();
   auto* mutableOffsets = offsets_->asMutable<vector_size_t>();
   auto* mutableSizes = sizes_->asMutable<vector_size_t>();
@@ -869,6 +916,12 @@ void ArrayVector::validate(const VectorValidateOptions& options) const {
   elements_->validate(options);
 }
 
+void ArrayVector::copyRanges(
+    const BaseVector* source,
+    const folly::Range<const CopyRange*>& ranges) {
+  copyRangesImpl(source, ranges, &elements_, nullptr);
+}
+
 bool MapVector::containsNullAt(vector_size_t idx) const {
   if (BaseVector::isNullAt(idx)) {
     return true;
@@ -1158,5 +1211,11 @@ void MapVector::validate(const VectorValidateOptions& options) const {
   values_->validate(options);
 }
 
+void MapVector::copyRanges(
+    const BaseVector* source,
+    const folly::Range<const CopyRange*>& ranges) {
+  copyRangesImpl(source, ranges, &values_, &keys_);
+}
+
 } // namespace velox
 } // namespace facebook

velox/vector/ComplexVector.h

@@ -317,9 +317,7 @@ struct ArrayVectorBase : BaseVector {
       const BaseVector* source,
       const folly::Range<const CopyRange*>& ranges,
       VectorPtr* targetValues,
-      const BaseVector* sourceValues,
-      VectorPtr* targetKeys,
-      const BaseVector* sourceKeys);
+      VectorPtr* targetKeys);
 
   void validateArrayVectorBase(
       const VectorValidateOptions& options,
@@ -407,20 +405,7 @@ class ArrayVector : public ArrayVectorBase {
 
   void copyRanges(
       const BaseVector* source,
-      const folly::Range<const CopyRange*>& ranges) override {
-    const ArrayVector* sourceArray{};
-    if (auto wrapped = source->wrappedVector();
-        !wrapped->isConstantEncoding()) {
-      sourceArray = wrapped->asUnchecked<ArrayVector>();
-    }
-    copyRangesImpl(
-        source,
-        ranges,
-        &elements_,
-        sourceArray ? sourceArray->elements_.get() : nullptr,
-        nullptr,
-        nullptr);
-  }
+      const folly::Range<const CopyRange*>& ranges) override;
 
   uint64_t retainedSize() const override {
     return BaseVector::retainedSize() + offsets_->capacity() +
@@ -547,20 +532,7 @@ class MapVector : public ArrayVectorBase {
 
   void copyRanges(
       const BaseVector* source,
-      const folly::Range<const CopyRange*>& ranges) override {
-    const MapVector* sourceMap{};
-    if (auto wrapped = source->wrappedVector();
-        !wrapped->isConstantEncoding()) {
-      sourceMap = wrapped->asUnchecked<MapVector>();
-    }
-    copyRangesImpl(
-        source,
-        ranges,
-        &values_,
-        sourceMap ? sourceMap->values_.get() : nullptr,
-        &keys_,
-        sourceMap ? sourceMap->keys_.get() : nullptr);
-  }
+      const folly::Range<const CopyRange*>& ranges) override;
 
   uint64_t retainedSize() const override {
     return BaseVector::retainedSize() + offsets_->capacity() +

velox/vector/FlatVector.cpp

@@ -51,6 +51,11 @@ void FlatVector<bool>::copyValuesAndNulls(
     const BaseVector* source,
     const SelectivityVector& rows,
     const vector_size_t* toSourceRow) {
+  if (source->typeKind() == TypeKind::UNKNOWN) {
+    rows.applyToSelected([&](auto row) { setNull(row, true); });
+    return;
+  }
+
   source = source->loadedVector();
   VELOX_CHECK(
       BaseVector::compatibleKind(BaseVector::typeKind(), source->typeKind()));
@@ -477,6 +482,15 @@ template <>
 void FlatVector<StringView>::copyRanges(
     const BaseVector* source,
     const folly::Range<const CopyRange*>& ranges) {
+  if (source->typeKind() == TypeKind::UNKNOWN) {
+    for (const auto& range : ranges) {
+      for (auto i = 0; i < range.count; ++i) {
+        setNull(range.targetIndex + i, true);
+      }
+    }
+    return;
+  }
+
   auto leaf = source->wrappedVector()->asUnchecked<SimpleVector<StringView>>();
   if (pool_ == leaf->pool()) {
     // We copy referencing the storage of 'source'.