refactor: 优化collectList函数

jetlinks · Dec 18, 2023 · adca411 · adca411
1 parent 9aa019a
commit adca411
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/src/main/java/org/jetlinks/reactor/ql/supports/agg/CollectListAggFeature.java b/src/main/java/org/jetlinks/reactor/ql/supports/agg/CollectListAggFeature.java
@@ -1,5 +1,6 @@
 package org.jetlinks.reactor.ql.supports.agg;
 
+import com.google.common.collect.Maps;
 import net.sf.jsqlparser.expression.Expression;
 import net.sf.jsqlparser.expression.StringValue;
 import net.sf.jsqlparser.schema.Column;
@@ -10,13 +11,12 @@
 import org.jetlinks.reactor.ql.ReactorQLRecord;
 import org.jetlinks.reactor.ql.feature.FeatureId;
 import org.jetlinks.reactor.ql.feature.FromFeature;
+import org.jetlinks.reactor.ql.feature.PropertyFeature;
 import org.jetlinks.reactor.ql.feature.ValueAggMapFeature;
 import reactor.core.publisher.Flux;
 
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -58,13 +58,20 @@ public Function<Flux<ReactorQLRecord>, Flux<Object>> createMapper(Expression exp
                     })
                     .collect(Collectors.toList());
 
+            PropertyFeature feature = metadata.getFeatureNow(PropertyFeature.ID);
+
             return flux -> flux
                     .map(record -> {
-                        Map<String, Object> values = new HashMap<>();
+                        Map<String, Object> values = Maps.newLinkedHashMapWithExpectedSize(columns.size());
+                        Map<String, Object> records = record.getRecords(true);
+                        Object row = record.getRecord();
                         for (String column : columns) {
-                            Optional.ofNullable(record.asMap())
-                                    .map(map -> map.get(column))
-                                    .ifPresent(val -> values.put(column, val));
+                            Object val = feature
+                                    .getProperty(column, records)
+                                    .orElseGet(() -> feature.getProperty(column, row).orElse(null));
+                            if (null != val) {
+                                values.put(column, val);
+                            }
                         }
                         return values;
                     })

diff --git a/src/test/java/org/jetlinks/reactor/ql/ReactorQLTest.java b/src/test/java/org/jetlinks/reactor/ql/ReactorQLTest.java
@@ -1428,6 +1428,23 @@ void testIfFunction() {
                  .verifyComplete();
     }
 
+    @Test
+    void testCollectList(){
+        String[] sql = {
+                "select collect_list(val) row from dual"
+        };
+        ReactorQL.builder()
+                 .sql(sql)
+                 .build()
+                 .start(Flux.range(1, 20)
+                            .map(i -> Collections.singletonMap("val", i)))
+                 .doOnNext(System.out::println)
+                 .as(StepVerifier::create)
+                 .expectNextCount(1)
+                 .verifyComplete();
+
+    }
+
     @Test
     void testCollectRowMap() {
         String[] sql = {