From 5e9f5502d3ad9da56f71ea6b2a455aa009020fff Mon Sep 17 00:00:00 2001 From: zhaorongsheng Date: Tue, 22 Sep 2026 11:42:03 +0800 Subject: [PATCH] [improvement](hive) Push partition filters to HMS ### What problem does this PR solve? Issue Number: #67724 Related PR: #67725 Problem Summary: Hive partition pruning previously enumerated every HMS partition before applying a selective predicate. The branch-4.2 architecture predates the master connector modules, so this change ports the direct HMS filter path to the legacy HMS table and scan pipeline. Plain Hive scans defer partition materialization, convert safe equality and IN predicates to the HMS filter grammar, and fetch only matching partitions through get_partitions_by_filter. Unsupported predicates, unsafe identifiers, saturated responses, and RPC failures retain the full-list local-pruning fallback; no-filter scans are materialized before translation. ### Release note Improve planning latency for selective Hive partition queries when HMS supports get_partitions_by_filter. ### Check List (For Author) - Test: Unit Test - HivePartitionFilterBuilderTest (3 tests), LogicalFileScanTest (10 tests), and PruneFileScanPartitionTest (1 test) - Maven offline FE test-compile for fe-core and dependent modules - Maven offline FE checkstyle validation - git diff --check - Behavior changed: Yes (selective plain-Hive scans can materialize only the HMS-filtered partition view) - Does this need documentation: No --- .../datasource/hive/HMSCachedClient.java | 8 + .../datasource/hive/HMSExternalTable.java | 69 +++++++ .../hive/HivePartitionFilterBuilder.java | 188 ++++++++++++++++++ .../hive/ThriftHMSCachedClient.java | 22 ++ .../translator/PhysicalPlanTranslator.java | 16 +- .../rules/rewrite/PruneFileScanPartition.java | 32 ++- .../trees/plans/logical/LogicalFileScan.java | 35 +++- .../hive/HivePartitionFilterBuilderTest.java | 80 ++++++++ .../plans/logical/LogicalFileScanTest.java | 18 ++ 9 files changed, 460 insertions(+), 8 deletions(-) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/datasource/hive/HivePartitionFilterBuilder.java create mode 100644 fe/fe-core/src/test/java/org/apache/doris/datasource/hive/HivePartitionFilterBuilderTest.java 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);