DuckLake: enforce partition predicates on temporal keys

docs/src/main/sphinx/connector/ducklake.md

@@ -233,6 +233,15 @@ The connector skips data files that cannot match query predicates:
   `file_statistics_pruning_enabled` [catalog session
   property](/sql/set-session) disable this behavior.
 
+A predicate on an `identity` partition key is enforced by the connector, not
+only used to prune. A `DELETE` that covers whole partitions of such a key ends
+those data files in the catalog, without reading a row of them. Keys of type
+`DATE`, `TIMESTAMP`, and `TIMESTAMP WITH TIME ZONE` are enforced along with the
+integer, boolean and `VARCHAR` keys. DuckDB writes some temporal values in a
+form the connector cannot read back — `infinity`, `-infinity`, BC dates, and
+years of five or more digits — and a table that holds one leaves the predicate
+to the engine, so the file is read and filtered instead of pruned.
+
 Before scheduling a data file, the coordinator reads its Parquet footer once
 and groups adjacent row groups into splits. `ducklake.max-split-size` is a soft
 target for the sum of the row groups' compressed column sizes. A row group is

plugin/trino-ducklake/src/main/java/io/trino/plugin/ducklake/DuckLakeMetadata.java

@@ -85,6 +85,8 @@
 import io.trino.spi.statistics.DoubleRange;
 import io.trino.spi.statistics.Estimate;
 import io.trino.spi.statistics.TableStatistics;
+import io.trino.spi.type.TimestampType;
+import io.trino.spi.type.TimestampWithTimeZoneType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.VarcharType;
 import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
@@ -124,6 +126,7 @@
 import static io.trino.spi.connector.RowChangeParadigm.DELETE_ROW_AND_INSERT_ROW;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
+import static io.trino.spi.type.DateType.DATE;
 import static io.trino.spi.type.DoubleType.DOUBLE;
 import static io.trino.spi.type.IntegerType.INTEGER;
 import static io.trino.spi.type.RealType.REAL;
@@ -1689,12 +1692,14 @@ public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(C
             newUnenforcedConstraint = TupleDomain.all();
         }
         else {
-            Set<Long> enforceableColumnIds = enforceablePartitionColumnIds(handle);
+            Map<Long, DuckLakePartitionColumn> partitionColumns = enforceablePartitionColumns(handle);
             ImmutableMap.Builder<DuckLakeColumnHandle, Domain> enforceableDomains = ImmutableMap.builder();
             ImmutableMap.Builder<DuckLakeColumnHandle, Domain> unenforceableDomains = ImmutableMap.builder();
             for (Map.Entry<DuckLakeColumnHandle, Domain> entry : predicate.getDomains().orElseThrow().entrySet()) {
                 DuckLakeColumnHandle column = entry.getKey();
-                if (enforceableColumnIds.contains(column.columnId()) && isEnforceableType(column.type())) {
+                DuckLakePartitionColumn partitionColumn = partitionColumns.get(column.columnId());
+                if (partitionColumn != null && isEnforceableType(column.type())
+                        && (!isTemporalType(column.type()) || everyPartitionValueDecidesTheColumn(handle, column, partitionColumn))) {
                     enforceableDomains.put(column, entry.getValue());
                 }
                 else {
@@ -1730,32 +1735,34 @@ public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(C
     }
 
     /**
-     * Returns the ids of partition columns whose predicates the connector can fully enforce by
-     * pruning files: identity-transformed columns where every visible data file was written with
-     * the current partitioning scheme, so every row of a kept file carries the partition value.
+     * Returns the partition columns whose predicates the connector can fully enforce by pruning
+     * files, by column id: identity-transformed columns where every visible data file was written
+     * with the current partitioning scheme, so every row of a kept file carries the partition
+     * value.
      */
-    private Set<Long> enforceablePartitionColumnIds(DuckLakeTableHandle handle)
+    private Map<Long, DuckLakePartitionColumn> enforceablePartitionColumns(DuckLakeTableHandle handle)
     {
         Optional<DuckLakePartitionInfo> partitionInfo = metastore.partitionInfo(handle.snapshotId(), handle.tableId());
         if (partitionInfo.isEmpty()) {
-            return ImmutableSet.of();
+            return ImmutableMap.of();
         }
         if (!metastore.allDataFilesUsePartition(handle.snapshotId(), handle.tableId(), partitionInfo.get().partitionId())) {
-            return ImmutableSet.of();
+            return ImmutableMap.of();
         }
         return partitionInfo.get().columns().stream()
                 .filter(column -> column.transform().equalsIgnoreCase(IDENTITY_TRANSFORM))
-                .map(DuckLakePartitionColumn::columnId)
-                .collect(toImmutableSet());
+                // a column filed under two identity keys is read from the first of them, which is
+                // the one PartitionTransforms picks as well
+                .collect(toImmutableMap(DuckLakePartitionColumn::columnId, column -> column, (first, _) -> first));
     }
 
     /**
-     * Types whose identity partition values round-trip exactly through the string
-     * representation in {@code ducklake_file_partition_value}, so pruning on them can be
-     * used to enforce a predicate. {@code DATE} is excluded because DuckDB writes values
-     * such as {@code infinity}, {@code -infinity}, BC dates and years with five or more
-     * digits that cannot be parsed back reliably; predicates on it stay unenforced and
-     * pruning remains fail-open.
+     * Types whose identity partition values round-trip exactly through the string representation
+     * in {@code ducklake_file_partition_value}, so pruning on them can be used to enforce a
+     * predicate. A temporal type is admitted here only as a candidate: DuckDB writes values such
+     * as {@code infinity}, {@code -infinity}, BC dates and years with five or more digits, which
+     * do not parse back, so {@link #isTemporalType} sends its columns through
+     * {@link #everyPartitionValueDecidesTheColumn} first.
      */
     private static boolean isEnforceableType(Type type)
     {
@@ -1764,7 +1771,40 @@ private static boolean isEnforceableType(Type type)
                 || type.equals(INTEGER)
                 || type.equals(BIGINT)
                 || type.equals(BOOLEAN)
-                || type instanceof VarcharType varcharType && varcharType.isUnbounded();
+                || type instanceof VarcharType varcharType && varcharType.isUnbounded()
+                || isTemporalType(type);
+    }
+
+    /**
+     * Types a partition value of which DuckDB may write in a form that does not parse back, so
+     * that the values of the table decide whether a predicate on the column can be enforced.
+     */
+    private static boolean isTemporalType(Type type)
+    {
+        return type.equals(DATE)
+                || type instanceof TimestampType
+                || type instanceof TimestampWithTimeZoneType;
+    }
+
+    /**
+     * Whether the partition value of every visible data file says which value of the column the
+     * rows of that file hold. A value the connector cannot read back leaves the file unprunable,
+     * and a predicate the connector cannot apply to one file it is not allowed to prune is one it
+     * cannot enforce for the table, so the engine keeps filtering the rows instead.
+     * <p>
+     * The values are read through the same code as {@link DuckLakeSplitManager#getSplits}, at the
+     * same snapshot, so a predicate enforced here is a predicate splits can be pruned by there.
+     * The query behind it costs one row per distinct value of the key, not one per data file, so
+     * a table of many files filed by few days is answered by few rows.
+     */
+    private boolean everyPartitionValueDecidesTheColumn(
+            DuckLakeTableHandle handle,
+            DuckLakeColumnHandle column,
+            DuckLakePartitionColumn partitionColumn)
+    {
+        List<DuckLakePartitionColumn> transforms = ImmutableList.of(partitionColumn);
+        return metastore.distinctPartitionValues(handle.snapshotId(), handle.tableId(), partitionColumn.partitionKeyIndex()).stream()
+                .allMatch(partitionValues -> PartitionTransforms.partitionDomain(column, transforms, partitionValues).isPresent());
     }
 
     /**

plugin/trino-ducklake/src/main/java/io/trino/plugin/ducklake/metastore/JdbcDuckLakeMetastore.java

@@ -646,6 +646,46 @@ SELECT count(*)
         }
     }
 
+    /**
+     * Returns the distinct partition values one key of the partitioning scheme takes across the
+     * data files of the table visible at the snapshot, each in the shape
+     * {@link DuckLakeDataFileEntry#partitionValues()} carries so that both are read the same way.
+     * A file that records no value under the key contributes the empty map.
+     * <p>
+     * The result holds one entry per distinct value rather than one per file, which is what makes
+     * it affordable to ask about a table of many files.
+     */
+    public List<Map<Integer, Optional<String>>> distinctPartitionValues(long snapshotId, long tableId, int partitionKeyIndex)
+    {
+        try (Handle handle = jdbi.open()) {
+            return handle.createQuery(
+                            """
+                            SELECT DISTINCT v.partition_key_index, v.partition_value
+                            FROM %s f
+                            LEFT JOIN %s v ON f.data_file_id = v.data_file_id AND f.table_id = v.table_id
+                                AND v.partition_key_index = :partitionKeyIndex
+                            WHERE f.table_id = :tableId AND %s""".formatted(
+                                    table("ducklake_data_file"),
+                                    table("ducklake_file_partition_value"),
+                                    visible("f")))
+                    .bind("snapshot", snapshotId)
+                    .bind("tableId", tableId)
+                    .bind("partitionKeyIndex", partitionKeyIndex)
+                    .map((rs, _) -> {
+                        int keyIndex = rs.getInt("partition_key_index");
+                        if (rs.wasNull()) {
+                            // a file that records nothing under the key says nothing about its rows
+                            return Map.<Integer, Optional<String>>of();
+                        }
+                        return Map.of(keyIndex, Optional.ofNullable(rs.getString("partition_value")));
+                    })
+                    .list();
+        }
+        catch (JdbiException e) {
+            throw metastoreError(e);
+        }
+    }
+
     public List<DuckLakeFileColumnStats> fileColumnStats(long tableId, Set<Long> columnIds)
     {
         if (columnIds.isEmpty()) {