This is an automated email from the ASF dual-hosted git repository.

huajianlan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git


The following commit(s) were added to refs/heads/master by this push:
     new 70f3e9d895d [chore](test) add some tests for prune nested column 
(#58354)
70f3e9d895d is described below

commit 70f3e9d895d4891ca9b27df953cb9dc9a96a47a7
Author: 924060929 <[email protected]>
AuthorDate: Wed Nov 26 10:35:50 2025 +0800

    [chore](test) add some tests for prune nested column (#58354)
    
    add some tests for prune nested column
---
 .../org/apache/doris/analysis/AccessPathInfo.java  | 45 ++++++++++--
 .../rewrite/AccessPathExpressionCollector.java     | 21 +++---
 .../nereids/rules/rewrite/NestedColumnPruning.java | 79 ++++++++++++++++------
 .../nereids/rules/rewrite/SlotTypeReplacer.java    |  7 +-
 .../rules/rewrite/PruneNestedColumnTest.java       | 76 +++++++++++++++++++++
 5 files changed, 191 insertions(+), 37 deletions(-)

diff --git 
a/fe/fe-core/src/main/java/org/apache/doris/analysis/AccessPathInfo.java 
b/fe/fe-core/src/main/java/org/apache/doris/analysis/AccessPathInfo.java
index 3df7baada4d..a2db1038a32 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/analysis/AccessPathInfo.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/analysis/AccessPathInfo.java
@@ -23,16 +23,51 @@ package org.apache.doris.analysis;
 import org.apache.doris.nereids.types.DataType;
 import org.apache.doris.thrift.TColumnAccessPath;
 
-import lombok.AllArgsConstructor;
-import lombok.Data;
-
 import java.util.List;
 
 /** AccessPathInfo */
-@Data
-@AllArgsConstructor
 public class AccessPathInfo {
+    public static final String ACCESS_ALL = "*";
+    public static final String ACCESS_MAP_KEYS = "KEYS";
+    public static final String ACCESS_MAP_VALUES = "VALUES";
+
     private DataType prunedType;
+    // allAccessPaths is used to record all access path include predicate 
access path and non-predicate access path,
+    // and predicateAccessPaths only contains the predicate access path.
+    // e.g. select element_at(s, 'name') from tbl where element_at(s, 'id') = 1
+    //      the allAccessPaths is: ["s.name", "s.id"]
+    //      the predicateAccessPaths is: ["s.id"]
     private List<TColumnAccessPath> allAccessPaths;
     private List<TColumnAccessPath> predicateAccessPaths;
+
+    public AccessPathInfo(DataType prunedType, List<TColumnAccessPath> 
allAccessPaths,
+            List<TColumnAccessPath> predicateAccessPaths) {
+        this.prunedType = prunedType;
+        this.allAccessPaths = allAccessPaths;
+        this.predicateAccessPaths = predicateAccessPaths;
+    }
+
+    public DataType getPrunedType() {
+        return prunedType;
+    }
+
+    public void setPrunedType(DataType prunedType) {
+        this.prunedType = prunedType;
+    }
+
+    public List<TColumnAccessPath> getAllAccessPaths() {
+        return allAccessPaths;
+    }
+
+    public void setAllAccessPaths(List<TColumnAccessPath> allAccessPaths) {
+        this.allAccessPaths = allAccessPaths;
+    }
+
+    public List<TColumnAccessPath> getPredicateAccessPaths() {
+        return predicateAccessPaths;
+    }
+
+    public void setPredicateAccessPaths(List<TColumnAccessPath> 
predicateAccessPaths) {
+        this.predicateAccessPaths = predicateAccessPaths;
+    }
 }
diff --git 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AccessPathExpressionCollector.java
 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AccessPathExpressionCollector.java
index 4fe9e9e2f74..9c665f379e9 100644
--- 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AccessPathExpressionCollector.java
+++ 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AccessPathExpressionCollector.java
@@ -17,6 +17,7 @@
 
 package org.apache.doris.nereids.rules.rewrite;
 
+import org.apache.doris.analysis.AccessPathInfo;
 import org.apache.doris.nereids.StatementContext;
 import 
org.apache.doris.nereids.rules.rewrite.AccessPathExpressionCollector.CollectorContext;
 import 
org.apache.doris.nereids.rules.rewrite.NestedColumnPruning.DataTypeAccessTree;
@@ -117,7 +118,7 @@ public class AccessPathExpressionCollector extends 
DefaultExpressionVisitor<Void
         if (nameToLambdaArguments.isEmpty()) {
             return null;
         }
-        context.accessPathBuilder.addPrefix("*");
+        context.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_ALL);
         Expression argument = 
nameToLambdaArguments.peek().get(arrayItemSlot.getName());
         if (argument == null) {
             return null;
@@ -157,7 +158,7 @@ public class AccessPathExpressionCollector extends 
DefaultExpressionVisitor<Void
         List<Expression> arguments = elementAt.getArguments();
         Expression first = arguments.get(0);
         if (first.getDataType().isArrayType() || 
first.getDataType().isMapType()) {
-            context.accessPathBuilder.addPrefix("*");
+            context.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_ALL);
             continueCollectAccessPath(first, context);
 
             for (int i = 1; i < arguments.size(); i++) {
@@ -200,39 +201,39 @@ public class AccessPathExpressionCollector extends 
DefaultExpressionVisitor<Void
     @Override
     public Void visitMapKeys(MapKeys mapKeys, CollectorContext context) {
         context = new CollectorContext(context.statementContext, 
context.bottomFilter);
-        context.accessPathBuilder.addPrefix("KEYS");
+        context.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_MAP_KEYS);
         return continueCollectAccessPath(mapKeys.getArgument(0), context);
     }
 
     @Override
     public Void visitMapValues(MapValues mapValues, CollectorContext context) {
         LinkedList<String> suffixPath = context.accessPathBuilder.accessPath;
-        if (!suffixPath.isEmpty() && suffixPath.get(0).equals("*")) {
+        if (!suffixPath.isEmpty() && 
suffixPath.get(0).equals(AccessPathInfo.ACCESS_ALL)) {
             CollectorContext removeStarContext
                     = new CollectorContext(context.statementContext, 
context.bottomFilter);
             
removeStarContext.accessPathBuilder.accessPath.addAll(suffixPath.subList(1, 
suffixPath.size()));
-            removeStarContext.accessPathBuilder.addPrefix("VALUES");
+            
removeStarContext.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_MAP_VALUES);
             return continueCollectAccessPath(mapValues.getArgument(0), 
removeStarContext);
         }
-        context.accessPathBuilder.addPrefix("VALUES");
+        context.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_MAP_VALUES);
         return continueCollectAccessPath(mapValues.getArgument(0), context);
     }
 
     @Override
     public Void visitMapContainsKey(MapContainsKey mapContainsKey, 
CollectorContext context) {
-        context.accessPathBuilder.addPrefix("KEYS");
+        context.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_MAP_KEYS);
         return continueCollectAccessPath(mapContainsKey.getArgument(0), 
context);
     }
 
     @Override
     public Void visitMapContainsValue(MapContainsValue mapContainsValue, 
CollectorContext context) {
-        context.accessPathBuilder.addPrefix("VALUES");
+        context.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_MAP_VALUES);
         return continueCollectAccessPath(mapContainsValue.getArgument(0), 
context);
     }
 
     @Override
     public Void visitMapContainsEntry(MapContainsEntry mapContainsEntry, 
CollectorContext context) {
-        context.accessPathBuilder.addPrefix("*");
+        context.accessPathBuilder.addPrefix(AccessPathInfo.ACCESS_ALL);
         return continueCollectAccessPath(mapContainsEntry.getArgument(0), 
context);
     }
 
@@ -398,7 +399,7 @@ public class AccessPathExpressionCollector extends 
DefaultExpressionVisitor<Void
         }
 
         List<String> path = context.accessPathBuilder.getPathList();
-        if (!path.isEmpty() && path.get(0).equals("*")) {
+        if (!path.isEmpty() && path.get(0).equals(AccessPathInfo.ACCESS_ALL)) {
             context.accessPathBuilder.removePrefix();
         }
 
diff --git 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/NestedColumnPruning.java
 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/NestedColumnPruning.java
index b03f183b115..016289ac139 100644
--- 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/NestedColumnPruning.java
+++ 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/NestedColumnPruning.java
@@ -244,11 +244,18 @@ public class NestedColumnPruning implements 
CustomRewriter {
 
     /** DataTypeAccessTree */
     public static class DataTypeAccessTree {
+        // type of this level
         private DataType type;
+        // is the root column?
         private boolean isRoot;
+        // if access 's.a.b' the node 's' and 'a' has accessPartialChild, and 
node 'b' has accessAll
         private boolean accessPartialChild;
         private boolean accessAll;
+        // for the future, only access the meta of the column,
+        // e.g. `is not null` can only access the column's offset, not need to 
read the data
         private TAccessPathType pathType;
+        // the children of the column, for example, column s is `struct<a:int, 
b:int>`,
+        // then node 's' has two children: 'a' and 'b', and the key is the 
column name
         private Map<String, DataTypeAccessTree> children = new 
LinkedHashMap<>();
 
         public DataTypeAccessTree(DataType type, TAccessPathType pathType) {
@@ -261,6 +268,30 @@ public class NestedColumnPruning implements CustomRewriter 
{
             this.pathType = pathType;
         }
 
+        public DataType getType() {
+            return type;
+        }
+
+        public boolean isRoot() {
+            return isRoot;
+        }
+
+        public boolean isAccessPartialChild() {
+            return accessPartialChild;
+        }
+
+        public boolean isAccessAll() {
+            return accessAll;
+        }
+
+        public TAccessPathType getPathType() {
+            return pathType;
+        }
+
+        public Map<String, DataTypeAccessTree> getChildren() {
+            return children;
+        }
+
         /** pruneCastType */
         public DataType pruneCastType(DataTypeAccessTree origin, 
DataTypeAccessTree cast) {
             if (type instanceof StructType) {
@@ -297,8 +328,16 @@ public class NestedColumnPruning implements CustomRewriter 
{
                 );
             } else if (type instanceof MapType) {
                 return MapType.of(
-                        
children.get("KEYS").pruneCastType(origin.children.get("KEYS"), 
cast.children.get("KEYS")),
-                        
children.get("VALUES").pruneCastType(origin.children.get("VALUES"), 
cast.children.get("VALUES"))
+                        children.get(AccessPathInfo.ACCESS_MAP_KEYS)
+                                .pruneCastType(
+                                        
origin.children.get(AccessPathInfo.ACCESS_MAP_KEYS),
+                                        
cast.children.get(AccessPathInfo.ACCESS_MAP_KEYS)
+                                ),
+                        children.get(AccessPathInfo.ACCESS_MAP_VALUES)
+                                .pruneCastType(
+                                        
origin.children.get(AccessPathInfo.ACCESS_MAP_VALUES),
+                                        
cast.children.get(AccessPathInfo.ACCESS_MAP_VALUES)
+                                )
                 );
             } else {
                 return cast.type;
@@ -327,7 +366,7 @@ public class NestedColumnPruning implements CustomRewriter {
                         cast.children.values().iterator().next(), path, index 
+ 1);
             } else if (cast.type instanceof MapType) {
                 String fieldName = path.get(index);
-                return children.get("VALUES").replacePathByAnotherTree(
+                return 
children.get(AccessPathInfo.ACCESS_MAP_VALUES).replacePathByAnotherTree(
                         cast.children.get(fieldName), path, index + 1
                 );
             }
@@ -358,33 +397,33 @@ public class NestedColumnPruning implements 
CustomRewriter {
                 }
                 return;
             } else if (this.type.isArrayType()) {
-                DataTypeAccessTree child = children.get("*");
-                if (path.get(accessIndex).equals("*")) {
+                DataTypeAccessTree child = 
children.get(AccessPathInfo.ACCESS_ALL);
+                if (path.get(accessIndex).equals(AccessPathInfo.ACCESS_ALL)) {
                     // enter this array and skip next *
                     child.setAccessByPath(path, accessIndex + 1, pathType);
                 }
                 return;
             } else if (this.type.isMapType()) {
                 String fieldName = path.get(accessIndex);
-                if (fieldName.equals("*")) {
+                if (fieldName.equals(AccessPathInfo.ACCESS_ALL)) {
                     // access value by the key, so we should access key and 
access value, then prune the value's type.
                     // e.g. map_column['id'] should access the keys, and 
access the values
-                    DataTypeAccessTree keysChild = children.get("KEYS");
-                    DataTypeAccessTree valuesChild = children.get("VALUES");
+                    DataTypeAccessTree keysChild = 
children.get(AccessPathInfo.ACCESS_MAP_KEYS);
+                    DataTypeAccessTree valuesChild = 
children.get(AccessPathInfo.ACCESS_MAP_VALUES);
                     keysChild.accessAll = true;
                     valuesChild.setAccessByPath(path, accessIndex + 1, 
pathType);
                     return;
-                } else if (fieldName.equals("KEYS")) {
+                } else if (fieldName.equals(AccessPathInfo.ACCESS_MAP_KEYS)) {
                     // only access the keys and not need enter keys, because 
it must be primitive type.
                     // e.g. map_keys(map_column)
-                    DataTypeAccessTree keysChild = children.get("KEYS");
+                    DataTypeAccessTree keysChild = 
children.get(AccessPathInfo.ACCESS_MAP_KEYS);
                     keysChild.accessAll = true;
                     return;
-                } else if (fieldName.equals("VALUES")) {
+                } else if (fieldName.equals(AccessPathInfo.ACCESS_MAP_VALUES)) 
{
                     // only access the values without keys, and maybe prune 
the value's data type.
                     // e.g. map_values(map_columns)[0] will access the array 
of values first,
                     //      and then access the array, so the access path is 
['VALUES', '*']
-                    DataTypeAccessTree valuesChild = children.get("VALUES");
+                    DataTypeAccessTree valuesChild = 
children.get(AccessPathInfo.ACCESS_MAP_VALUES);
                     valuesChild.setAccessByPath(path, accessIndex + 1, 
pathType);
                     return;
                 }
@@ -411,10 +450,10 @@ public class NestedColumnPruning implements 
CustomRewriter {
                     root.children.put(kv.getKey().toLowerCase(), 
of(kv.getValue().getDataType(), pathType));
                 }
             } else if (type instanceof ArrayType) {
-                root.children.put("*", of(((ArrayType) type).getItemType(), 
pathType));
+                root.children.put(AccessPathInfo.ACCESS_ALL, of(((ArrayType) 
type).getItemType(), pathType));
             } else if (type instanceof MapType) {
-                root.children.put("KEYS", of(((MapType) type).getKeyType(), 
pathType));
-                root.children.put("VALUES", of(((MapType) 
type).getValueType(), pathType));
+                root.children.put(AccessPathInfo.ACCESS_MAP_KEYS, 
of(((MapType) type).getKeyType(), pathType));
+                root.children.put(AccessPathInfo.ACCESS_MAP_VALUES, 
of(((MapType) type).getValueType(), pathType));
             }
             return root;
         }
@@ -440,17 +479,17 @@ public class NestedColumnPruning implements 
CustomRewriter {
                     }
                 }
             } else if (type instanceof ArrayType) {
-                Optional<DataType> childDataType = 
children.get("*").pruneDataType();
+                Optional<DataType> childDataType = 
children.get(AccessPathInfo.ACCESS_ALL).pruneDataType();
                 if (childDataType.isPresent()) {
-                    accessedChildren.add(Pair.of("*", childDataType.get()));
+                    accessedChildren.add(Pair.of(AccessPathInfo.ACCESS_ALL, 
childDataType.get()));
                 }
             } else if (type instanceof MapType) {
-                DataType prunedValueType = children.get("VALUES")
+                DataType prunedValueType = 
children.get(AccessPathInfo.ACCESS_MAP_VALUES)
                         .pruneDataType()
                         .orElse(((MapType) type).getValueType());
                 // can not prune keys but can prune values
-                accessedChildren.add(Pair.of("KEYS", ((MapType) 
type).getKeyType()));
-                accessedChildren.add(Pair.of("VALUES", prunedValueType));
+                accessedChildren.add(Pair.of(AccessPathInfo.ACCESS_MAP_KEYS, 
((MapType) type).getKeyType()));
+                accessedChildren.add(Pair.of(AccessPathInfo.ACCESS_MAP_VALUES, 
prunedValueType));
             }
             if (accessedChildren.isEmpty()) {
                 return Optional.of(type);
diff --git 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SlotTypeReplacer.java
 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SlotTypeReplacer.java
index 9f6a02391eb..661f537eabf 100644
--- 
a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SlotTypeReplacer.java
+++ 
b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/SlotTypeReplacer.java
@@ -21,6 +21,7 @@ import org.apache.doris.analysis.AccessPathInfo;
 import org.apache.doris.catalog.Column;
 import org.apache.doris.common.Pair;
 import org.apache.doris.datasource.iceberg.IcebergExternalTable;
+import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.properties.OrderKey;
 import 
org.apache.doris.nereids.rules.rewrite.NestedColumnPruning.DataTypeAccessTree;
 import org.apache.doris.nereids.trees.expressions.ArrayItemReference;
@@ -542,8 +543,10 @@ public class SlotTypeReplacer extends 
DefaultPlanRewriter<Void> {
         ImmutableCollection.Builder<E> newExprs;
         if (expressions instanceof List) {
             newExprs = ImmutableList.builder();
-        } else {
+        } else if (expressions instanceof Set) {
             newExprs = ImmutableSet.builder();
+        } else {
+            throw new AnalysisException("Unsupported expression type: " + 
expressions.getClass());
         }
 
         boolean changed = false;
@@ -688,7 +691,7 @@ public class SlotTypeReplacer extends 
DefaultPlanRewriter<Void> {
                         originPath, index + 1, ((ArrayType) 
type).getItemType(), column.getChildren().get(0)
                 );
             } else if (type instanceof MapType) {
-                if (fieldName.equals("*") || fieldName.equals("VALUES")) {
+                if (fieldName.equals(AccessPathInfo.ACCESS_ALL) || 
fieldName.equals(AccessPathInfo.ACCESS_MAP_VALUES)) {
                     replaceIcebergAccessPathToId(
                             originPath, index + 1, ((MapType) 
type).getValueType(), column.getChildren().get(1)
                     );
diff --git 
a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java
 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java
index 9c6674deb3d..22e9c5e949f 100644
--- 
a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java
+++ 
b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java
@@ -23,6 +23,7 @@ import org.apache.doris.common.Pair;
 import org.apache.doris.common.Triple;
 import org.apache.doris.nereids.NereidsPlanner;
 import org.apache.doris.nereids.rules.RuleType;
+import 
org.apache.doris.nereids.rules.rewrite.NestedColumnPruning.DataTypeAccessTree;
 import org.apache.doris.nereids.trees.expressions.Alias;
 import org.apache.doris.nereids.trees.expressions.ArrayItemReference;
 import org.apache.doris.nereids.trees.expressions.Expression;
@@ -31,10 +32,14 @@ import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.SlotReference;
 import 
org.apache.doris.nereids.trees.expressions.functions.scalar.StructElement;
 import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
+import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalCTEConsumer;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalUnion;
 import org.apache.doris.nereids.types.DataType;
+import org.apache.doris.nereids.types.NestedColumnPrunable;
+import org.apache.doris.nereids.types.NullType;
 import org.apache.doris.nereids.util.MemoPatternMatchSupported;
 import org.apache.doris.nereids.util.PlanChecker;
 import org.apache.doris.planner.OlapScanNode;
@@ -537,6 +542,77 @@ public class PruneNestedColumnTest extends 
TestWithFeService implements MemoPatt
                 );
     }
 
+    @Test
+    public void testDataTypeAccessTree() {
+        List<Pair<SlotReference, DataTypeAccessTree>> trees = 
getDataTypeAccessTrees(
+                "select struct_element(s, 'city') from (select id, s from tbl 
union all select 1, null) tmp");
+
+        Assertions.assertEquals(1, trees.size());
+        DataTypeAccessTree tree = trees.get(0).second;
+        Assertions.assertEquals(NullType.INSTANCE, tree.getType());
+        Assertions.assertEquals(1, tree.getChildren().size());
+        Assertions.assertEquals("STRUCT<city:TEXT>", 
tree.getChildren().get("s").getType().toSql());
+
+        SlotReference slot = trees.get(0).first;
+        Type columnType = slot.getOriginalColumn().get().getType();
+        
Assertions.assertEquals("struct<city:text,data:array<map<int,struct<a:int,b:double>>>>",
 columnType.toSql());
+
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "city"), 
"STRUCT<city:TEXT>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data"), 
"STRUCT<data:ARRAY<MAP<INT,STRUCT<a:INT,b:DOUBLE>>>>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data", "*"), 
"STRUCT<data:ARRAY<MAP<INT,STRUCT<a:INT,b:DOUBLE>>>>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data", "*", 
"KEYS"), "STRUCT<data:ARRAY<MAP<INT,STRUCT<a:INT,b:DOUBLE>>>>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data", "*", 
"VALUES"), "STRUCT<data:ARRAY<MAP<INT,STRUCT<a:INT,b:DOUBLE>>>>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data", "*", 
"VALUES", "a"), "STRUCT<data:ARRAY<MAP<INT,STRUCT<a:INT>>>>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data", "*", 
"VALUES", "b"), "STRUCT<data:ARRAY<MAP<INT,STRUCT<b:DOUBLE>>>>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data", "*", 
"*"), "STRUCT<data:ARRAY<MAP<INT,STRUCT<a:INT,b:DOUBLE>>>>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data", "*", 
"*", "a"), "STRUCT<data:ARRAY<MAP<INT,STRUCT<a:INT>>>>");
+        setAccessPathAndAssertType(slot, ImmutableList.of("s", "data", "*", 
"*", "b"), "STRUCT<data:ARRAY<MAP<INT,STRUCT<b:DOUBLE>>>>");
+
+        setAccessPathsAndAssertType(slot,
+                ImmutableList.of(
+                        ImmutableList.of("s", "data", "*", "*", "b"),
+                        ImmutableList.of("s", "city")
+                ),
+                "STRUCT<city:TEXT,data:ARRAY<MAP<INT,STRUCT<b:DOUBLE>>>>"
+        );
+    }
+
+    private void setAccessPathAndAssertType(SlotReference slot, List<String> 
path, String expectedType) {
+        setAccessPathsAndAssertType(slot, ImmutableList.of(path), 
expectedType);
+    }
+
+    private void setAccessPathsAndAssertType(SlotReference slot, 
List<List<String>> paths, String expectedType) {
+        DataType columnType = 
DataType.fromCatalogType(slot.getOriginalColumn().get().getType());
+        SlotReference originColumnTypeSlot = new SlotReference(slot.getName(), 
columnType);
+        DataTypeAccessTree tree = 
DataTypeAccessTree.ofRoot(originColumnTypeSlot, TAccessPathType.DATA);
+        for (List<String> path : paths) {
+            tree.setAccessByPath(path, 0, TAccessPathType.DATA);
+        }
+        DataType dataType = tree.pruneDataType().get();
+        Assertions.assertEquals(expectedType, dataType.toSql());
+    }
+
+    private List<Pair<SlotReference, DataTypeAccessTree>> 
getDataTypeAccessTrees(String sql) {
+        Plan rewritePlan = PlanChecker.from(connectContext)
+                .analyze(sql)
+                .rewrite()
+                .getCascadesContext()
+                .getRewritePlan();
+
+        List<Slot> output = ((LogicalOlapScan) 
rewritePlan.collect(LogicalOlapScan.class::isInstance)
+                .iterator()
+                .next())
+                .getOutput();
+        List<Pair<SlotReference, DataTypeAccessTree>> trees = new 
ArrayList<>();
+        for (Slot slot : output) {
+            if (slot.getDataType() instanceof NestedColumnPrunable) {
+                DataTypeAccessTree dataTypeAccessTree = 
DataTypeAccessTree.ofRoot(slot, TAccessPathType.DATA);
+                trees.add(Pair.of((SlotReference) slot, dataTypeAccessTree));
+            }
+        }
+        return trees;
+    }
+
     private void assertColumn(String sql, String expectType,
             List<TColumnAccessPath> expectAllAccessPaths,
             List<TColumnAccessPath> expectPredicateAccessPaths) throws 
Exception {


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to