Fix partition predicate pushdown containing null in Glue

trinodb · Jul 8, 2022 · 2a75a34 · 2a75a34
1 parent 432d289
commit 2a75a34
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 0 deletions.
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueExpressionUtil.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueExpressionUtil.java
@@ -196,6 +196,10 @@ else if (singleValues.size() > 1) {
             disjuncts.add(inClause);
         }
 
+        if (domain.isNullAllowed()) {
+            disjuncts.add(format("(%s = '%s')", columnName, NULL_STRING));
+        }
+
         return Optional.of("(" + DISJUNCT_JOINER.join(disjuncts) + ")");
     }
 }
diff --git a/.../trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueExpressionUtil.java b/.../trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestGlueExpressionUtil.java
@@ -144,6 +144,28 @@ public void testBuildGlueExpressionTupleDomainNotNull()
         assertEquals(expression, format("(col1 <> '%s')", GlueExpressionUtil.NULL_STRING));
     }
 
+    @Test
+    public void testBuildGlueExpressionTupleDomainEqualsOrIsNull()
+    {
+        TupleDomain<String> filter = new PartitionFilterBuilder()
+                .addStringValues("col1", "2020-01-01")
+                .addDomain("col1", Domain.onlyNull(VarcharType.VARCHAR))
+                .build();
+        String expression = buildGlueExpression(ImmutableList.of("col1"), filter, true);
+        assertEquals(expression, format("((col1 = '2020-01-01') OR (col1 = '%s'))", GlueExpressionUtil.NULL_STRING));
+    }
+
+    @Test
+    public void testBuildGlueExpressionTupleDomainEqualsAndIsNotNull()
+    {
+        TupleDomain<String> filter = new PartitionFilterBuilder()
+                .addStringValues("col1", "2020-01-01")
+                .addDomain("col2", Domain.notNull(VarcharType.VARCHAR))
+                .build();
+        String expression = buildGlueExpression(ImmutableList.of("col1", "col2"), filter, true);
+        assertEquals(expression, format("((col1 = '2020-01-01')) AND (col2 <> '%s')", GlueExpressionUtil.NULL_STRING));
+    }
+
     @Test
     public void testBuildGlueExpressionMaxLengthNone()
     {

diff --git a/...n/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java b/...n/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java
@@ -837,6 +837,49 @@ public void testGetPartitionsFilterIsNullWithValue()
                 ImmutableList.of(ImmutableList.of(GlueExpressionUtil.NULL_STRING)));
     }
 
+    @Test
+    public void testGetPartitionsFilterEqualsOrIsNullWithValue()
+            throws Exception
+    {
+        TupleDomain<String> equalsOrIsNullFilter = new PartitionFilterBuilder()
+                .addStringValues(PARTITION_KEY, "2020-03-01")
+                .addDomain(PARTITION_KEY, Domain.onlyNull(VarcharType.VARCHAR))
+                .build();
+        List<String> partitionList = new ArrayList<>();
+        partitionList.add("2020-01-01");
+        partitionList.add("2020-02-01");
+        partitionList.add("2020-03-01");
+        partitionList.add(null);
+
+        doGetPartitionsFilterTest(
+                CREATE_TABLE_COLUMNS_PARTITIONED_VARCHAR,
+                PARTITION_KEY,
+                partitionList,
+                ImmutableList.of(equalsOrIsNullFilter),
+                ImmutableList.of(ImmutableList.of("2020-03-01", GlueExpressionUtil.NULL_STRING)));
+    }
+
+    @Test
+    public void testGetPartitionsFilterEqualsAndIsNotNull()
+            throws Exception
+    {
+        TupleDomain<String> equalsAndIsNotNullFilter = new PartitionFilterBuilder()
+                .addDomain(PARTITION_KEY, Domain.notNull(VarcharType.VARCHAR))
+                .addBigintValues(PARTITION_KEY2, 300L)
+                .build();
+
+        doGetPartitionsFilterTest(
+                CREATE_TABLE_COLUMNS_PARTITIONED_TWO_KEYS,
+                ImmutableList.of(PARTITION_KEY, PARTITION_KEY2),
+                ImmutableList.of(
+                        PartitionValues.make("2020-01-01", "100"),
+                        PartitionValues.make("2020-02-01", "200"),
+                        PartitionValues.make("2020-03-01", "300"),
+                        PartitionValues.make(null, "300")),
+                ImmutableList.of(equalsAndIsNotNullFilter),
+                ImmutableList.of(ImmutableList.of(PartitionValues.make("2020-03-01", "300"))));
+    }
+
     @Test
     public void testUpdateStatisticsOnCreate()
     {