From 8ea17ca082cf2215a05942e23eec38c70d5a0aff Mon Sep 17 00:00:00 2001
From: Jimmy Lu <jimmylu@meta.com>
Date: Fri, 26 Apr 2024 11:23:52 -0700
Subject: [PATCH] Fix crash when map subfield is pruned but we still want to
 add non-null filter on its keys

Differential Revision: D56642855
---
 velox/connectors/hive/HiveConnectorUtil.cpp   |  7 ++++---
 .../hive/tests/HiveConnectorTest.cpp          | 19 +++++++++++++++++++
 velox/dwio/common/ScanSpec.h                  |  4 ++++
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/velox/connectors/hive/HiveConnectorUtil.cpp b/velox/connectors/hive/HiveConnectorUtil.cpp
index ecdf6d2c48c2..2fadb1fd8acd 100644
--- a/velox/connectors/hive/HiveConnectorUtil.cpp
+++ b/velox/connectors/hive/HiveConnectorUtil.cpp
@@ -323,9 +323,10 @@ namespace {
 
 void filterOutNullMapKeys(const Type& rootType, common::ScanSpec& rootSpec) {
   rootSpec.visit(rootType, [](const Type& type, common::ScanSpec& spec) {
-    if (type.isMap()) {
-      spec.childByName(common::ScanSpec::kMapKeysFieldName)
-          ->addFilter(common::IsNotNull());
+    if (type.isMap() && !spec.isConstant()) {
+      auto* keys = spec.childByName(common::ScanSpec::kMapKeysFieldName);
+      VELOX_CHECK_NOT_NULL(keys);
+      keys->addFilter(common::IsNotNull());
     }
   });
 }
diff --git a/velox/connectors/hive/tests/HiveConnectorTest.cpp b/velox/connectors/hive/tests/HiveConnectorTest.cpp
index 3daa1d792dff..ebc0afdcb6b4 100644
--- a/velox/connectors/hive/tests/HiveConnectorTest.cpp
+++ b/velox/connectors/hive/tests/HiveConnectorTest.cpp
@@ -417,6 +417,25 @@ TEST_F(HiveConnectorTest, makeScanSpec_filterPartitionKey) {
   ASSERT_FALSE(scanSpec->childByName("ds")->projectOut());
 }
 
+TEST_F(HiveConnectorTest, makeScanSpec_prunedMapNonNullMapKey) {
+  auto rowType =
+      ROW({"c0"},
+          {ROW(
+              {{"c0c0", MAP(BIGINT(), MAP(BIGINT(), BIGINT()))},
+               {"c0c1", BIGINT()}})});
+  auto scanSpec = makeScanSpec(
+      rowType,
+      groupSubfields(makeSubfields({"c0.c0c1"})),
+      {},
+      nullptr,
+      {},
+      {},
+      pool_.get());
+  auto* c0 = scanSpec->childByName("c0");
+  ASSERT_EQ(c0->children().size(), 2);
+  ASSERT_TRUE(c0->childByName("c0c0")->isConstant());
+}
+
 TEST_F(HiveConnectorTest, extractFiltersFromRemainingFilter) {
   core::QueryCtx queryCtx;
   exec::SimpleExpressionEvaluator evaluator(&queryCtx, pool_.get());
diff --git a/velox/dwio/common/ScanSpec.h b/velox/dwio/common/ScanSpec.h
index 9d998cce896a..75ad6e3935e3 100644
--- a/velox/dwio/common/ScanSpec.h
+++ b/velox/dwio/common/ScanSpec.h
@@ -410,6 +410,10 @@ class ScanSpec {
 template <typename F>
 void ScanSpec::visit(const Type& type, F&& f) {
   f(type, *this);
+  if (isConstant()) {
+    // Child specs are not populated in this case.
+    return;
+  }
   switch (type.kind()) {
     case TypeKind::ROW:
       for (auto& child : children_) {