Filter hive splits on $path constraints

Earlier splits were only pruned if the $path predicate was
representable as a Domain

Author: raunaqmorarka

plugin/trino-hive/src/main/java/io/trino/plugin/hive/BackgroundHiveSplitLoader.java

@@ -52,6 +52,7 @@
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.connector.Constraint;
 import io.trino.spi.connector.DynamicFilter;
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.spi.type.TypeManager;
@@ -148,6 +149,7 @@ public class BackgroundHiveSplitLoader
 
     private final Table table;
     private final TupleDomain<? extends ColumnHandle> compactEffectivePredicate;
+    private final Constraint constraint;
     private final DynamicFilter dynamicFilter;
     private final long dynamicFilteringWaitTimeoutMillis;
     private final TypeManager typeManager;
@@ -192,6 +194,7 @@ public BackgroundHiveSplitLoader(
             Table table,
             Iterator<HivePartitionMetadata> partitions,
             TupleDomain<? extends ColumnHandle> compactEffectivePredicate,
+            Constraint constraint,
             DynamicFilter dynamicFilter,
             Duration dynamicFilteringWaitTimeout,
             TypeManager typeManager,
@@ -209,6 +212,7 @@ public BackgroundHiveSplitLoader(
     {
         this.table = table;
         this.compactEffectivePredicate = compactEffectivePredicate;
+        this.constraint = constraint;
         this.dynamicFilter = dynamicFilter;
         this.dynamicFilteringWaitTimeoutMillis = dynamicFilteringWaitTimeout.toMillis();
         this.typeManager = typeManager;
@@ -423,6 +427,7 @@ private ListenableFuture<Void> loadPartition(HivePartitionMetadata partition)
                     schema,
                     partitionKeys,
                     effectivePredicate,
+                    constraint,
                     partitionMatchSupplier,
                     partition.getHiveColumnCoercions(),
                     Optional.empty(),
@@ -475,6 +480,7 @@ private ListenableFuture<Void> loadPartition(HivePartitionMetadata partition)
                 schema,
                 partitionKeys,
                 effectivePredicate,
+                constraint,
                 partitionMatchSupplier,
                 partition.getHiveColumnCoercions(),
                 bucketConversionRequiresWorkerParticipation ? bucketConversion : Optional.empty(),

plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveSplitManager.java

@@ -263,6 +263,7 @@ public ConnectorSplitSource getSplits(
                 table,
                 hivePartitions,
                 hiveTable.getCompactEffectivePredicate(),
+                constraint,
                 dynamicFilter,
                 getDynamicFilteringWaitTimeout(session),
                 typeManager,

plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/InternalHiveSplitFactory.java

@@ -15,6 +15,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import io.airlift.slice.Slice;
 import io.airlift.units.DataSize;
 import io.trino.metastore.HiveTypeName;
 import io.trino.plugin.hive.AcidInfo;
@@ -32,7 +33,10 @@
 import io.trino.plugin.hive.parquet.ParquetPageSourceFactory;
 import io.trino.plugin.hive.rcfile.RcFilePageSourceFactory;
 import io.trino.spi.HostAddress;
+import io.trino.spi.connector.ColumnHandle;
+import io.trino.spi.connector.Constraint;
 import io.trino.spi.predicate.Domain;
+import io.trino.spi.predicate.NullableValue;
 import io.trino.spi.predicate.TupleDomain;
 
 import java.util.Collection;
@@ -41,11 +45,14 @@
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.function.BooleanSupplier;
+import java.util.function.Predicate;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.airlift.slice.Slices.utf8Slice;
+import static io.trino.plugin.hive.HiveColumnHandle.PATH_TYPE;
 import static io.trino.plugin.hive.HiveColumnHandle.isPathColumnHandle;
+import static io.trino.plugin.hive.HiveColumnHandle.pathColumnHandle;
 import static io.trino.plugin.hive.util.AcidTables.isFullAcidTable;
 import static io.trino.plugin.hive.util.HiveUtil.getSerializationLibraryName;
 import static java.util.Objects.requireNonNull;
@@ -58,6 +65,7 @@ public class InternalHiveSplitFactory
     private final List<HivePartitionKey> partitionKeys;
     private final Optional<Domain> pathDomain;
     private final Map<Integer, HiveTypeName> hiveColumnCoercions;
+    private final Constraint constraint;
     private final BooleanSupplier partitionMatchSupplier;
     private final Optional<BucketConversion> bucketConversion;
     private final Optional<HiveSplit.BucketValidation> bucketValidation;
@@ -71,6 +79,7 @@ public InternalHiveSplitFactory(
             Map<String, String> schema,
             List<HivePartitionKey> partitionKeys,
             TupleDomain<HiveColumnHandle> effectivePredicate,
+            Constraint constraint,
             BooleanSupplier partitionMatchSupplier,
             Map<Integer, HiveTypeName> hiveColumnCoercions,
             Optional<BucketConversion> bucketConversion,
@@ -84,6 +93,7 @@ public InternalHiveSplitFactory(
         this.strippedSchema = stripUnnecessaryProperties(requireNonNull(schema, "schema is null"));
         this.partitionKeys = requireNonNull(partitionKeys, "partitionKeys is null");
         pathDomain = getPathDomain(requireNonNull(effectivePredicate, "effectivePredicate is null"));
+        this.constraint = requireNonNull(constraint, "constraint is null");
         this.partitionMatchSupplier = requireNonNull(partitionMatchSupplier, "partitionMatchSupplier is null");
         this.hiveColumnCoercions = ImmutableMap.copyOf(requireNonNull(hiveColumnCoercions, "hiveColumnCoercions is null"));
         this.bucketConversion = requireNonNull(bucketConversion, "bucketConversion is null");
@@ -144,7 +154,7 @@ private Optional<InternalHiveSplit> createInternalHiveSplit(
             boolean splittable,
             Optional<AcidInfo> acidInfo)
     {
-        if (!pathMatchesPredicate(pathDomain, path)) {
+        if (!pathMatchesPredicate(pathDomain, constraint, path)) {
             return Optional.empty();
         }
 
@@ -259,10 +269,16 @@ private static Optional<Domain> getPathDomain(TupleDomain<HiveColumnHandle> effe
                         .findFirst());
     }
 
-    private static boolean pathMatchesPredicate(Optional<Domain> pathDomain, String path)
+    private static boolean pathMatchesPredicate(Optional<Domain> pathDomain, Constraint constraint, String path)
     {
-        return pathDomain
-                .map(domain -> domain.includesNullableValue(utf8Slice(path)))
-                .orElse(true);
+        Slice pathSlice = utf8Slice(path);
+        if (pathDomain.isPresent() && !pathDomain.get().includesNullableValue(pathSlice)) {
+            return false;
+        }
+        if (constraint.getPredicateColumns().isEmpty() || !constraint.getPredicateColumns().get().contains(pathColumnHandle())) {
+            return true;
+        }
+        Predicate<Map<ColumnHandle, NullableValue>> predicate = constraint.predicate().orElse(_ -> true);
+        return predicate.test(ImmutableMap.of(pathColumnHandle(), new NullableValue(PATH_TYPE, pathSlice)));
     }
 }

plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseHiveConnectorTest.java

@@ -5599,6 +5599,33 @@ public void testBucketFilteringByInPredicate()
         }
     }
 
+    @Test
+    void testPathConstraintSplitPruning()
+    {
+        String tableName = "test_path_constraint_split_pruning_" + randomNameSuffix();
+        assertUpdate("CREATE TABLE " + tableName + " (id bigint, data varchar)");
+        assertUpdate("INSERT INTO " + tableName + " VALUES (1, 'a'), (2, 'b'), (3, 'c')", 3);
+        assertUpdate("INSERT INTO " + tableName + " VALUES (1, 'a'), (2, 'b'), (3, 'c')", 3);
+
+        try {
+            Set<String> files = getTableFiles(tableName);
+            assertThat(files).hasSizeGreaterThanOrEqualTo(2);
+            String fileName = files.stream().findAny().orElseThrow();
+            String query = "SELECT * FROM " + tableName + " WHERE \"$path\" LIKE '%" + fileName + "%'";
+            assertQueryStats(
+                    getSession(),
+                    query,
+                    queryStats -> {
+                        assertThat(queryStats.getTotalDrivers()).isEqualTo(1);
+                        assertThat(queryStats.getPhysicalInputPositions()).isEqualTo(3);
+                    },
+                    results -> assertThat(results.getRowCount()).isEqualTo(3));
+        }
+        finally {
+            assertUpdate("DROP TABLE " + tableName);
+        }
+    }
+
     @Test
     public void testSchemaMismatchesWithDereferenceProjections()
     {

plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestBackgroundHiveSplitLoader.java

@@ -47,6 +47,7 @@
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorSplitSource.ConnectorSplitBatch;
+import io.trino.spi.connector.Constraint;
 import io.trino.spi.connector.DynamicFilter;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.predicate.Domain;
@@ -154,7 +155,7 @@ public void tearDown()
     public void testNoPathFilter()
             throws Exception
     {
-        BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(TEST_LOCATIONS, TupleDomain.none());
+        BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(TEST_LOCATIONS, TupleDomain.none(), Constraint.alwaysTrue());
 
         HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader);
         backgroundHiveSplitLoader.start(hiveSplitSource);
@@ -205,7 +206,8 @@ public void testPathFilter()
     {
         BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(
                 TEST_LOCATIONS,
-                LOCATION_DOMAIN);
+                LOCATION_DOMAIN,
+                Constraint.alwaysTrue());
 
         HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader);
         backgroundHiveSplitLoader.start(hiveSplitSource);
@@ -214,13 +216,32 @@ public void testPathFilter()
         assertThat(paths.get(0)).isEqualTo(LOCATION.toString());
     }
 
+    @Test
+    public void testPathConstraint()
+            throws Exception
+    {
+        BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(
+                TEST_LOCATIONS,
+                LOCATION_DOMAIN,
+                new Constraint(
+                        TupleDomain.all(),
+                        LOCATION_DOMAIN.transformKeys(ColumnHandle.class::cast).asPredicate(),
+                        Set.of(pathColumnHandle())));
+
+        HiveSplitSource hiveSplitSource = hiveSplitSource(backgroundHiveSplitLoader);
+        backgroundHiveSplitLoader.start(hiveSplitSource);
+        List<String> paths = drain(hiveSplitSource);
+        assertThat(paths).containsExactly(LOCATION.toString());
+    }
+
     @Test
     public void testPathFilterOneBucketMatchPartitionedTable()
             throws Exception
     {
         BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(
                 TEST_LOCATIONS,
                 LOCATION_DOMAIN,
+                Constraint.alwaysTrue(),
                 Optional.of(new HiveBucketFilter(Set.of(0, 1))),
                 PARTITIONED_TABLE,
                 Optional.of(new HiveTablePartitioning(true, BUCKETING_V1, BUCKET_COUNT, BUCKET_COLUMN_HANDLES, false, List.of(), true)));
@@ -239,6 +260,7 @@ public void testPathFilterBucketedPartitionedTable()
         BackgroundHiveSplitLoader backgroundHiveSplitLoader = backgroundHiveSplitLoader(
                 TEST_LOCATIONS,
                 LOCATION_DOMAIN,
+                Constraint.alwaysTrue(),
                 Optional.empty(),
                 PARTITIONED_TABLE,
                 Optional.of(
@@ -474,6 +496,7 @@ public HivePartitionMetadata next()
                     }
                 },
                 TupleDomain.all(),
+                Constraint.alwaysTrue(),
                 DynamicFilter.EMPTY,
                 new Duration(0, SECONDS),
                 TESTING_TYPE_MANAGER,
@@ -796,6 +819,7 @@ public void testBuildManifestFileIterator()
                 schema,
                 List.of(),
                 TupleDomain.all(),
+                Constraint.alwaysTrue(),
                 () -> true,
                 ImmutableMap.of(),
                 Optional.empty(),
@@ -837,6 +861,7 @@ public void testBuildManifestFileIteratorNestedDirectory()
                 schema,
                 List.of(),
                 TupleDomain.all(),
+                Constraint.alwaysTrue(),
                 () -> true,
                 ImmutableMap.of(),
                 Optional.empty(),
@@ -875,6 +900,7 @@ public void testBuildManifestFileIteratorWithCacheInvalidation()
                 schema,
                 List.of(),
                 TupleDomain.all(),
+                Constraint.alwaysTrue(),
                 () -> true,
                 ImmutableMap.of(),
                 Optional.empty(),
@@ -1056,6 +1082,7 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
         return backgroundHiveSplitLoader(
                 fileSystemFactory,
                 TupleDomain.all(),
+                Constraint.alwaysTrue(),
                 dynamicFilter,
                 dynamicFilteringProbeBlockingTimeoutMillis,
                 Optional.empty(),
@@ -1079,12 +1106,14 @@ private static TrinoFileSystemFactory createTestingFileSystem(Collection<Locatio
 
     private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
             List<Location> locations,
-            TupleDomain<HiveColumnHandle> tupleDomain)
+            TupleDomain<HiveColumnHandle> tupleDomain,
+            Constraint constraint)
             throws IOException
     {
         return backgroundHiveSplitLoader(
                 locations,
                 tupleDomain,
+                constraint,
                 Optional.empty(),
                 SIMPLE_TABLE,
                 Optional.empty());
@@ -1093,37 +1122,23 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
     private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
             List<Location> locations,
             TupleDomain<HiveColumnHandle> compactEffectivePredicate,
+            Constraint constraint,
             Optional<HiveBucketFilter> hiveBucketFilter,
             Table table,
             Optional<HiveTablePartitioning> tablePartitioning)
             throws IOException
-    {
-        return backgroundHiveSplitLoader(
-                locations,
-                compactEffectivePredicate,
-                hiveBucketFilter,
-                table,
-                tablePartitioning,
-                Optional.empty());
-    }
-
-    private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
-            List<Location> locations,
-            TupleDomain<HiveColumnHandle> compactEffectivePredicate,
-            Optional<HiveBucketFilter> hiveBucketFilter,
-            Table table,
-            Optional<HiveTablePartitioning> tablePartitioning,
-            Optional<ValidWriteIdList> validWriteIds)
-            throws IOException
     {
         TrinoFileSystemFactory fileSystemFactory = createTestingFileSystem(locations);
         return backgroundHiveSplitLoader(
                 fileSystemFactory,
                 compactEffectivePredicate,
+                constraint,
+                DynamicFilter.EMPTY,
+                new Duration(0, SECONDS),
                 hiveBucketFilter,
                 table,
                 tablePartitioning,
-                validWriteIds);
+                Optional.empty());
     }
 
     private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
@@ -1137,6 +1152,7 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
         return backgroundHiveSplitLoader(
                 fileSystemFactory,
                 compactEffectivePredicate,
+                Constraint.alwaysTrue(),
                 DynamicFilter.EMPTY,
                 new Duration(0, SECONDS),
                 hiveBucketFilter,
@@ -1148,6 +1164,7 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
     private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
             TrinoFileSystemFactory fileSystemFactory,
             TupleDomain<HiveColumnHandle> compactEffectivePredicate,
+            Constraint constraint,
             DynamicFilter dynamicFilter,
             Duration dynamicFilteringProbeBlockingTimeout,
             Optional<HiveBucketFilter> hiveBucketFilter,
@@ -1166,6 +1183,7 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
                 table,
                 hivePartitionMetadatas.iterator(),
                 compactEffectivePredicate,
+                constraint,
                 dynamicFilter,
                 dynamicFilteringProbeBlockingTimeout,
                 TESTING_TYPE_MANAGER,
@@ -1210,6 +1228,7 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoader(
                 SIMPLE_TABLE,
                 partitions.iterator(),
                 TupleDomain.none(),
+                Constraint.alwaysTrue(),
                 DynamicFilter.EMPTY,
                 new Duration(0, SECONDS),
                 TESTING_TYPE_MANAGER,
@@ -1237,6 +1256,7 @@ private BackgroundHiveSplitLoader backgroundHiveSplitLoaderOfflinePartitions()
                 SIMPLE_TABLE,
                 createPartitionMetadataWithOfflinePartitions(),
                 TupleDomain.all(),
+                Constraint.alwaysTrue(),
                 DynamicFilter.EMPTY,
                 new Duration(0, SECONDS),
                 TESTING_TYPE_MANAGER,