diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HMSCachedClient.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HMSCachedClient.java index a3a04600c9fdee..9a3b631899ea27 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HMSCachedClient.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HMSCachedClient.java @@ -53,6 +53,14 @@ public interface HMSCachedClient { List listPartitions(String dbName, String tblName); + /** + * Lists partitions matching the HMS filter grammar. Callers use this to avoid enumerating every + * partition before applying a selective equality or IN predicate. + */ + default List listPartitionsByFilter(String dbName, String tblName, String filter) { + throw new UnsupportedOperationException("listPartitionsByFilter is not supported"); + } + List listPartitionNames(String dbName, String tblName, long maxListPartitionNum); Partition getPartition(String dbName, String tblName, List partitionValues); diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HMSExternalTable.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HMSExternalTable.java index 2a7759faecb735..d03693b31fbb59 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HMSExternalTable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HMSExternalTable.java @@ -17,6 +17,7 @@ package org.apache.doris.datasource.hive; +import org.apache.doris.analysis.PartitionValue; import org.apache.doris.analysis.TableScanParams; import org.apache.doris.analysis.TableSnapshot; import org.apache.doris.catalog.Column; @@ -24,6 +25,7 @@ import org.apache.doris.catalog.ListPartitionItem; import org.apache.doris.catalog.MTMV; import org.apache.doris.catalog.PartitionItem; +import org.apache.doris.catalog.PartitionKey; import org.apache.doris.catalog.PartitionType; import org.apache.doris.catalog.PrimitiveType; import org.apache.doris.catalog.ScalarType; @@ -62,6 +64,7 @@ import org.apache.doris.mtmv.MTMVSnapshotIf; import org.apache.doris.nereids.exceptions.NotSupportedException; import org.apache.doris.nereids.rules.expression.rules.SortedPartitionRanges; +import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.plans.algebra.CatalogRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan.SelectedPartitions; import org.apache.doris.qe.GlobalVariable; @@ -97,6 +100,7 @@ import org.apache.hadoop.hive.metastore.api.Partition; import org.apache.hadoop.hive.metastore.api.StringColumnStatsData; import org.apache.hadoop.hive.metastore.api.Table; +import org.apache.hadoop.hive.metastore.utils.FileUtils; import org.apache.hadoop.hive.ql.io.AcidUtils; import org.apache.hudi.common.table.HoodieTableMetaClient; import org.apache.hudi.internal.schema.InternalSchema; @@ -467,6 +471,42 @@ public Optional> getSortedPartitionRanges(CatalogR return hivePartitionValues.getSortedPartitionRanges(); } + @Override + public SelectedPartitions initSelectedPartitions(Optional snapshot) { + if (getDlaType() == DLAType.HIVE && !getPartitionColumns(snapshot).isEmpty()) { + return SelectedPartitions.DEFERRED_PARTITION_PRUNING; + } + return super.initSelectedPartitions(snapshot); + } + + /** + * Materializes only the partitions admitted by a safe HMS partition filter. An empty result means + * the predicate or metastore does not support the filter grammar and local pruning remains intact. + */ + public Optional> getNameToPartitionItemsByFilter( + Optional snapshot, Expression predicate) { + if (getDlaType() != DLAType.HIVE) { + return Optional.empty(); + } + List partitionColumns = getPartitionColumns(snapshot); + if (partitionColumns.isEmpty()) { + return Optional.empty(); + } + String filter = HivePartitionFilterBuilder.build(predicate, partitionColumns); + if (filter == null) { + return Optional.empty(); + } + try { + List partitions = ((HMSExternalCatalog) catalog).getClient() + .listPartitionsByFilter(getRemoteDbName(), getRemoteName(), filter); + return Optional.of(toNameToPartitionItems(partitions, partitionColumns, snapshot)); + } catch (RuntimeException e) { + LOG.warn("Failed to prune Hive partitions through HMS filter for table {}.{}", + getDbName(), getName()); + return Optional.empty(); + } + } + public SelectedPartitions initHudiSelectedPartitions(Optional tableSnapshot) { if (getDlaType() != DLAType.HUDI) { return SelectedPartitions.NOT_PRUNED; @@ -518,6 +558,35 @@ public Map getNameToPartitionItems() { return nameToPartitionItem; } + private Map toNameToPartitionItems(List partitions, + List partitionColumns, Optional snapshot) { + List partitionColumnNames = partitionColumns.stream() + .map(Column::getName) + .collect(Collectors.toList()); + List partitionColumnTypes = getPartitionColumnTypes(snapshot); + Map result = Maps.newHashMapWithExpectedSize(partitions.size()); + for (Partition partition : partitions) { + String partitionName = FileUtils.makePartName(partitionColumnNames, partition.getValues()); + result.put(partitionName, toListPartitionItem(partitionName, partitionColumnTypes)); + } + return result; + } + + private ListPartitionItem toListPartitionItem(String partitionName, List types) { + List partitionValues = HiveUtil.toPartitionValues(partitionName); + List values = Lists.newArrayListWithExpectedSize(types.size()); + for (String partitionValue : partitionValues) { + values.add(new PartitionValue(partitionValue, + HiveExternalMetaCache.HIVE_DEFAULT_PARTITION.equals(partitionValue))); + } + try { + PartitionKey partitionKey = PartitionKey.createListPartitionKeyWithTypes(values, types, true); + return new ListPartitionItem(Lists.newArrayList(partitionKey)); + } catch (AnalysisException e) { + throw new HMSClientException("failed to convert filtered Hive partition %s", e, partitionName); + } + } + public boolean isHiveTransactionalTable() { return dlaType == DLAType.HIVE && AcidUtils.isTransactionalTable(remoteTable); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HivePartitionFilterBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HivePartitionFilterBuilder.java new file mode 100644 index 00000000000000..bac196e3ed0f04 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HivePartitionFilterBuilder.java @@ -0,0 +1,188 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource.hive; + +import org.apache.doris.catalog.Column; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.InPredicate; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.types.DataType; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.Lists; + +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** Converts the safe subset of Nereids partition predicates to the HMS filter grammar. */ +public final class HivePartitionFilterBuilder { + + private HivePartitionFilterBuilder() { + } + + /** + * Builds an HMS filter from direct equality and IN predicates. Returns null when any predicate + * or column type is outside the supported grammar, so callers can retain local pruning. + */ + public static String build(Expression predicate, List partitionColumns) { + Map columnsByName = partitionColumns.stream() + .collect(Collectors.toMap(column -> column.getName().toLowerCase(Locale.ROOT), + Function.identity())); + Map> valuesByName = new HashMap<>(); + for (Expression conjunct : ExpressionUtils.extractConjunction(predicate)) { + if (!collectValues(conjunct, columnsByName, valuesByName)) { + return null; + } + } + return buildFilter(partitionColumns, valuesByName); + } + + private static boolean collectValues(Expression expression, Map columnsByName, + Map> valuesByName) { + if (expression instanceof EqualTo) { + EqualTo equalTo = (EqualTo) expression; + return collectEqualValue(equalTo.left(), equalTo.right(), columnsByName, valuesByName) + || collectEqualValue(equalTo.right(), equalTo.left(), columnsByName, valuesByName); + } + if (expression instanceof InPredicate) { + InPredicate inPredicate = (InPredicate) expression; + String columnName = slotName(inPredicate.getCompareExpr()); + if (columnName == null || !columnsByName.containsKey(columnName)) { + return false; + } + List values = Lists.newArrayList(); + for (Expression option : inPredicate.getOptions()) { + String value = literalValue(option); + if (value == null + || !literalIsSupported((Literal) option, columnsByName.get(columnName))) { + return false; + } + values.add(value); + } + valuesByName.computeIfAbsent(columnName, ignored -> Lists.newArrayList()).addAll(values); + return true; + } + return false; + } + + private static boolean collectEqualValue(Expression slotExpression, Expression literalExpression, + Map columnsByName, Map> valuesByName) { + String columnName = slotName(slotExpression); + if (columnName == null || !columnsByName.containsKey(columnName)) { + return false; + } + String value = literalValue(literalExpression); + if (value == null + || !literalIsSupported((Literal) literalExpression, columnsByName.get(columnName))) { + return false; + } + valuesByName.computeIfAbsent(columnName, ignored -> Lists.newArrayList()).add(value); + return true; + } + + private static String slotName(Expression expression) { + if (expression instanceof SlotReference) { + return ((SlotReference) expression).getName().toLowerCase(Locale.ROOT); + } + return null; + } + + private static String literalValue(Expression expression) { + if (!(expression instanceof Literal)) { + return null; + } + Object value = ((Literal) expression).getValue(); + return value == null ? null : value.toString(); + } + + private static boolean literalIsSupported(Literal literal, Column column) { + DataType literalType = literal.getDataType(); + String value = literal.getValue().toString(); + if (column.getType().isIntegerType()) { + return literalType.isIntegerType() && isIntegralLiteral(literal.getValue().toString()); + } + return column.getType().isStringType() && literalType.isStringType() + && value.indexOf('\\') < 0 && value.indexOf('\'') < 0; + } + + private static String buildFilter(List partitionColumns, Map> valuesByName) { + List filters = Lists.newArrayList(); + for (Column partitionColumn : partitionColumns) { + List values = valuesByName.get(partitionColumn.getName().toLowerCase(Locale.ROOT)); + if (values == null || values.isEmpty()) { + continue; + } + if (!isHmsFilterIdentifier(partitionColumn.getName())) { + return null; + } + List valueFilters = values.stream() + .map(value -> partitionColumn.getName() + " = " + toHmsLiteral(value, partitionColumn)) + .collect(Collectors.toList()); + filters.add("(" + String.join(" OR ", valueFilters) + ")"); + } + return filters.isEmpty() ? null : String.join(" AND ", filters); + } + + private static String toHmsLiteral(String value, Column column) { + return column.getType().isIntegerType() ? value : "'" + value + "'"; + } + + private static boolean isHmsFilterIdentifier(String value) { + if (value.isEmpty() || !isHmsFilterLetterOrDigit(value.charAt(0))) { + return false; + } + String lowerCaseValue = value.toLowerCase(Locale.ROOT); + if (lowerCaseValue.equals("not") || lowerCaseValue.equals("and") || lowerCaseValue.equals("or") + || lowerCaseValue.equals("like") || lowerCaseValue.equals("date") + || lowerCaseValue.equals("const") || lowerCaseValue.equals("struct") + || isAllDigits(value)) { + return false; + } + for (int index = 1; index < value.length(); index++) { + char character = value.charAt(index); + if (!isHmsFilterLetterOrDigit(character) && character != '_') { + return false; + } + } + return true; + } + + private static boolean isAllDigits(String value) { + return value.chars().allMatch(character -> character >= '0' && character <= '9'); + } + + private static boolean isHmsFilterLetterOrDigit(char value) { + return value >= 'a' && value <= 'z' + || value >= 'A' && value <= 'Z' + || value >= '0' && value <= '9'; + } + + private static boolean isIntegralLiteral(String value) { + int start = value.startsWith("-") ? 1 : 0; + if (start == value.length()) { + return false; + } + return value.chars().skip(start).allMatch(character -> character >= '0' && character <= '9'); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/ThriftHMSCachedClient.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/ThriftHMSCachedClient.java index 0ddb0e0f48ef3c..b35f6f7c291d20 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/ThriftHMSCachedClient.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/hive/ThriftHMSCachedClient.java @@ -332,6 +332,28 @@ public List listPartitions(String dbName, String tblName) { } } + @Override + public List listPartitionsByFilter(String dbName, String tblName, String filter) { + short maxPartitions = (short) (DEFAULT_PARTITION_BATCH_SIZE + 1); + try (ThriftHMSClient client = getClient()) { + try { + List partitions = ugiDoAs(() -> client.client.listPartitionsByFilter( + dbName, tblName, filter, maxPartitions)); + if (partitions.size() > DEFAULT_PARTITION_BATCH_SIZE) { + throw new HMSClientException( + "HMS partition filter matched more than %d partitions in table '%s.%s'.", + DEFAULT_PARTITION_BATCH_SIZE, dbName, tblName); + } + return partitions; + } catch (Exception e) { + client.setThrowable(e); + throw e; + } + } catch (Exception e) { + throw new HMSClientException("failed to filter partitions in table '%s.%s'.", e, dbName, tblName); + } + } + @Override public List listPartitionNames(String dbName, String tblName, long maxListPartitionNum) { // list all parts when the limit is greater than the short maximum diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java index 59508d1779311a..0733a53b2f0746 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java @@ -43,6 +43,7 @@ import org.apache.doris.catalog.Function.NullableMode; import org.apache.doris.catalog.OdbcTable; import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.PartitionItem; import org.apache.doris.catalog.TableIf; import org.apache.doris.catalog.Type; import org.apache.doris.common.Config; @@ -127,6 +128,7 @@ import org.apache.doris.nereids.trees.plans.PreAggStatus; import org.apache.doris.nereids.trees.plans.algebra.Aggregate; import org.apache.doris.nereids.trees.plans.algebra.Relation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFileScan; import org.apache.doris.nereids.trees.plans.physical.AbstractPhysicalJoin; import org.apache.doris.nereids.trees.plans.physical.AbstractPhysicalSort; import org.apache.doris.nereids.trees.plans.physical.PhysicalAssertNumRows; @@ -747,7 +749,8 @@ public PlanFragment visitPhysicalFileScan(PhysicalFileScan fileScan, PlanTransla scanNode = new HiveScanNode(context.nextPlanNodeId(), tupleDescriptor, false, sv, directoryLister, context.getScanContext()); HiveScanNode hiveScanNode = (HiveScanNode) scanNode; - hiveScanNode.setSelectedPartitions(fileScan.getSelectedPartitions()); + hiveScanNode.setSelectedPartitions( + materializeDeferredHivePartitions((HMSExternalTable) table, fileScan)); if (fileScan.getTableSample().isPresent()) { hiveScanNode.setTableSample(new TableSample(fileScan.getTableSample().get().isPercent, fileScan.getTableSample().get().sampleValue, fileScan.getTableSample().get().seek)); @@ -796,6 +799,17 @@ public PlanFragment visitPhysicalFileScan(PhysicalFileScan fileScan, PlanTransla return getPlanFragmentForPhysicalFileScan(fileScan, context, scanNode, table, tupleDescriptor); } + private LogicalFileScan.SelectedPartitions materializeDeferredHivePartitions( + HMSExternalTable table, PhysicalFileScan fileScan) { + LogicalFileScan.SelectedPartitions selectedPartitions = fileScan.getSelectedPartitions(); + if (!selectedPartitions.isDeferredPartitionPruning()) { + return selectedPartitions; + } + Map partitionItems = + table.getNameToPartitionItems(fileScan.getRelationSnapshot()); + return new LogicalFileScan.SelectedPartitions(partitionItems.size(), partitionItems, false); + } + @Override public PlanFragment visitPhysicalEmptyRelation(PhysicalEmptyRelation emptyRelation, PlanTranslatorContext context) { List output = emptyRelation.getOutput(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneFileScanPartition.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneFileScanPartition.java index 3022fc7d7e8d27..e77214bb8876aa 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneFileScanPartition.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PruneFileScanPartition.java @@ -20,6 +20,7 @@ import org.apache.doris.catalog.Column; import org.apache.doris.catalog.PartitionItem; import org.apache.doris.datasource.ExternalTable; +import org.apache.doris.datasource.hive.HMSExternalTable; import org.apache.doris.nereids.CascadesContext; import org.apache.doris.nereids.rules.Rule; import org.apache.doris.nereids.rules.RuleType; @@ -91,22 +92,43 @@ private SelectedPartitions pruneExternalPartitions(ExternalTable externalTable, .collect(Collectors.toList()); Map nameToPartitionItem = scan.getSelectedPartitions().selectedPartitions; + boolean connectorFilteredPartitions = false; + if (nameToPartitionItem.isEmpty() + && scan.getSelectedPartitions().isDeferredPartitionPruning() + && externalTable instanceof HMSExternalTable) { + Optional> filteredPartitions = + ((HMSExternalTable) externalTable).getNameToPartitionItemsByFilter( + scan.getRelationSnapshot(), filter.getPredicate()); + if (filteredPartitions.isPresent()) { + nameToPartitionItem = filteredPartitions.get(); + connectorFilteredPartitions = true; + } + } + if (!connectorFilteredPartitions && nameToPartitionItem.isEmpty() + && (scan.getSelectedPartitions().isNotPruned() + || scan.getSelectedPartitions().isDeferredPartitionPruning())) { + nameToPartitionItem = externalTable.getNameToPartitionItems(scan.getRelationSnapshot()); + } + final Map partitionItems = nameToPartitionItem; Optional> sortedPartitionRanges = Optional.empty(); boolean enableBinarySearch = ctx.getConnectContext() == null || ctx.getConnectContext().getSessionVariable().enableBinarySearchFilteringPartitions; if (enableBinarySearch) { - sortedPartitionRanges = (Optional) externalTable.getSortedPartitionRanges(scan); + sortedPartitionRanges = connectorFilteredPartitions + ? Optional.ofNullable(SortedPartitionRanges.build(partitionItems)) + : (Optional) externalTable.getSortedPartitionRanges(scan); } PartitionPruneResult result = PartitionPruner.pruneWithResult( - partitionSlots, filter.getPredicate(), nameToPartitionItem, ctx, + partitionSlots, filter.getPredicate(), partitionItems, ctx, PartitionTableType.EXTERNAL, sortedPartitionRanges); List prunedPartitions = new ArrayList<>(result.partitions); for (String name : prunedPartitions) { - selectedPartitionItems.put(name, nameToPartitionItem.get(name)); + selectedPartitionItems.put(name, partitionItems.get(name)); } - return new SelectedPartitions(nameToPartitionItem.size(), selectedPartitionItems, true, - result.hasPartitionPredicate); + return new SelectedPartitions( + connectorFilteredPartitions ? -1L : partitionItems.size(), + selectedPartitionItems, true, connectorFilteredPartitions || result.hasPartitionPredicate); } static List getPartitionColumnsForScan(ExternalTable externalTable, LogicalFileScan scan) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalFileScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalFileScan.java index 12dfffea883152..967826d3d70a25 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalFileScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalFileScan.java @@ -383,10 +383,24 @@ private boolean hasSameScanParams(Optional left, Optional sel */ public SelectedPartitions(long totalPartitionNum, Map selectedPartitions, boolean isPruned, boolean hasPartitionPredicate) { + this(totalPartitionNum, selectedPartitions, isPruned, hasPartitionPredicate, State.MATERIALIZED); + } + + private SelectedPartitions(long totalPartitionNum, Map selectedPartitions, + boolean isPruned, boolean hasPartitionPredicate, State state) { this.totalPartitionNum = totalPartitionNum; this.selectedPartitions = ImmutableMap.copyOf(Objects.requireNonNull(selectedPartitions, "selectedPartitions is null")); this.isPruned = isPruned; this.hasPartitionPredicate = hasPartitionPredicate; + this.state = state; + } + + public boolean isNotPruned() { + return state == State.NOT_PRUNED; + } + + public boolean isDeferredPartitionPruning() { + return state == State.DEFERRED; } @Override @@ -437,13 +467,14 @@ public boolean equals(Object o) { SelectedPartitions that = (SelectedPartitions) o; return isPruned == that.isPruned && hasPartitionPredicate == that.hasPartitionPredicate + && state == that.state && Objects.equals( selectedPartitions.keySet(), that.selectedPartitions.keySet()); } @Override public int hashCode() { - return Objects.hash(selectedPartitions, isPruned, hasPartitionPredicate); + return Objects.hash(selectedPartitions, isPruned, hasPartitionPredicate, state); } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/HivePartitionFilterBuilderTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/HivePartitionFilterBuilderTest.java new file mode 100644 index 00000000000000..a63ddf99d4a6d5 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/hive/HivePartitionFilterBuilderTest.java @@ -0,0 +1,80 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.datasource.hive; + +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.Type; +import org.apache.doris.nereids.trees.expressions.And; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.InPredicate; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.types.StringType; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.List; + +public class HivePartitionFilterBuilderTest { + @Test + public void testBuildsEqualityAndInFilter() { + Column city = new Column("city", Type.STRING, true); + Column day = new Column("day", Type.INT, true); + SlotReference citySlot = new SlotReference("city", StringType.INSTANCE); + SlotReference daySlot = new SlotReference("day", IntegerType.INSTANCE); + List predicates = Arrays.asList( + new EqualTo(citySlot, new StringLiteral("shanghai")), + new InPredicate(daySlot, Arrays.asList(new IntegerLiteral(1), new IntegerLiteral(2)))); + + Assertions.assertEquals("(city = 'shanghai') AND (day = 1 OR day = 2)", + HivePartitionFilterBuilder.build(new And(predicates), Arrays.asList(city, day))); + } + + @Test + public void testRejectsUnsafeStringAndUnsupportedPredicate() { + Column city = new Column("city", Type.STRING, true); + Column day = new Column("day", Type.INT, true); + SlotReference citySlot = new SlotReference("city", StringType.INSTANCE); + SlotReference daySlot = new SlotReference("day", IntegerType.INSTANCE); + + Assertions.assertNull(HivePartitionFilterBuilder.build( + new EqualTo(citySlot, new StringLiteral("can't")), Arrays.asList(city))); + Assertions.assertNull(HivePartitionFilterBuilder.build( + new EqualTo(daySlot, new StringLiteral("1")), Arrays.asList(day))); + Assertions.assertNull(HivePartitionFilterBuilder.build( + new EqualTo(citySlot, daySlot), + Arrays.asList(city))); + } + + @Test + public void testRejectsFilterUnsafeColumnNames() { + Column date = new Column("date", Type.STRING, true); + Column digits = new Column("123", Type.STRING, true); + SlotReference dateSlot = new SlotReference("date", StringType.INSTANCE); + SlotReference digitsSlot = new SlotReference("123", StringType.INSTANCE); + + Assertions.assertNull(HivePartitionFilterBuilder.build( + new EqualTo(dateSlot, new StringLiteral("20260101")), Arrays.asList(date))); + Assertions.assertNull(HivePartitionFilterBuilder.build( + new EqualTo(digitsSlot, new StringLiteral("x")), Arrays.asList(digits))); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalFileScanTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalFileScanTest.java index eaaa24ebaf0f03..d72fcda5f994b5 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalFileScanTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/logical/LogicalFileScanTest.java @@ -48,6 +48,24 @@ public class LogicalFileScanTest { + @Test + public void testDeferredSelectedPartitionsState() { + HMSExternalTable table = Mockito.mock(HMSExternalTable.class); + Mockito.when(table.initSelectedPartitions(Mockito.any())) + .thenReturn(SelectedPartitions.DEFERRED_PARTITION_PRUNING); + Mockito.when(table.getFullSchema(Mockito.any())) + .thenReturn(Collections.singletonList(new Column("id", Type.INT, true))); + Mockito.when(table.getName()).thenReturn("hive_tbl"); + + LogicalFileScan scan = new LogicalFileScan(new RelationId(12), table, + Collections.singletonList("db"), Collections.emptyList(), + Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty()); + + Assertions.assertSame(SelectedPartitions.DEFERRED_PARTITION_PRUNING, scan.getSelectedPartitions()); + Assertions.assertTrue(scan.getSelectedPartitions().isDeferredPartitionPruning()); + Assertions.assertFalse(scan.getSelectedPartitions().isNotPruned()); + } + @Test public void testComputeOutputIncludesInvisibleRowLineageColumnsForIcebergTable() { Column rowIdColumn = new Column(IcebergUtils.ICEBERG_ROW_ID_COL, Type.BIGINT, true);