From 97e4f0c11e18dbfde505e5b82143e0bd9512e0e4 Mon Sep 17 00:00:00 2001
From: Vladimir Sitnikov <sitnikov.vladimir@gmail.com>
Date: Wed, 2 Sep 2020 23:54:59 +0300
Subject: [PATCH] [CALCITE-4215] Use default trait value if traitSupplier
 returns null

---
 .../main/java/org/apache/calcite/plan/RelTraitSet.java    | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java
index e6737dc7f99f..8e4766a5a156 100644
--- a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java
+++ b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java
@@ -240,6 +240,9 @@ public <T extends RelMultipleTrait> RelTraitSet replaceIfs(RelTraitDef<T> def,
       return this; // trait is not enabled; ignore it
     }
     final List<T> traitList = traitSupplier.get();
+    if (traitList == null) {
+      return replace(index, def.getDefault());
+    }
     return replace(index, RelCompositeTrait.of(def, traitList));
   }
 
@@ -250,7 +253,10 @@ public <T extends RelTrait> RelTraitSet replaceIf(RelTraitDef<T> def,
     if (index < 0) {
       return this; // trait is not enabled; ignore it
     }
-    final T traitList = traitSupplier.get();
+    T traitList = traitSupplier.get();
+    if (traitList == null) {
+      traitList = def.getDefault();
+    }
     return replace(index, traitList);
   }