Fix CASE expression caches leaking across TsBlocks

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBCaseWhenThenTableIT.java

@@ -695,4 +695,54 @@ public void testThenWithBinarySameConstant() {
     String[] retArray = new String[] {"0,", "11,", "22,", "33,"};
     tableResultSetEqualTest(sql, expectedHeader, retArray, DATABASE);
   }
+
+  @Test
+  public void testCaseWhenReusesAggregationAcrossBlocks() {
+    prepareTableData(
+        new String[] {
+          "USE " + DATABASE,
+          "CREATE TABLE case_aggregation_cache (dev STRING TAG, c DOUBLE FIELD, q INT32 FIELD)",
+          "INSERT INTO case_aggregation_cache(time, dev, c, q) VALUES "
+              + "(0, 'd0', 1.0, 0), (0, 'd1', 1.0, 0), (3600000, 'd1', 1.0, 0)"
+        });
+
+    // The non-pushdown predicate and the different device/window shapes make the aggregation
+    // produce multiple blocks. CASE clears the shared aggregate before the next output uses it.
+    String suffix = " FROM case_aggregation_cache WHERE q % 2 < 1 GROUP BY dev, date_bin(1h, time)";
+    for (String projection :
+        new String[] {
+          "CASE WHEN sum(c) > 0 THEN 1.0 END AS case_result, CAST(sum(c) AS DOUBLE) AS total",
+          "CASE WHEN sum(c) > 0 THEN 1.0 END AS case_result, sum(c) + 0 AS total",
+          "CASE WHEN sum(c) > 0 THEN CAST(sum(c) AS DOUBLE) / CAST(sum(c) AS DOUBLE) END "
+              + "AS case_result, CAST(sum(c) AS DOUBLE) AS total"
+        }) {
+      tableResultSetEqualTest(
+          "SELECT " + projection + suffix,
+          new String[] {"case_result", "total"},
+          new String[] {"1.0,1.0,", "1.0,1.0,", "1.0,1.0,"},
+          DATABASE);
+    }
+  }
+
+  @Test
+  public void testReusingCasePreservesPendingSiblingColumn() {
+    prepareTableData(
+        new String[] {
+          "USE " + DATABASE,
+          "CREATE TABLE case_sibling_cache (dev STRING TAG, x DOUBLE FIELD)",
+          "INSERT INTO case_sibling_cache(time, dev, x) VALUES "
+              + "(0, 'd0', -3.0), (1, 'd0', 0.0), (2, 'd0', 2.0)"
+        });
+
+    String sharedCase = "CASE WHEN x > 0 THEN 1.0 ELSE 0.0 END";
+    tableResultSetEqualTest(
+        "SELECT "
+            + sharedCase
+            + " AS a, x + CASE WHEN ("
+            + sharedCase
+            + ") > 0 THEN 2.0 ELSE 3.0 END AS b FROM case_sibling_cache ORDER BY time",
+        new String[] {"a", "b"},
+        new String[] {"0.0,0.0,", "0.0,3.0,", "1.0,4.0,"},
+        DATABASE);
+  }
 }

integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBCaseDeviceCacheIT.java

@@ -0,0 +1,136 @@
+/*
+ * 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.iotdb.relational.it.schema;
+
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.TableClusterIT;
+import org.apache.iotdb.itbase.category.TableLocalStandaloneIT;
+import org.apache.iotdb.itbase.env.BaseEnv;
+
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+
+import static org.apache.iotdb.db.it.utils.TestUtils.assertResultSetEqual;
+import static org.apache.iotdb.db.it.utils.TestUtils.prepareTableData;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({TableLocalStandaloneIT.class, TableClusterIT.class})
+public class IoTDBCaseDeviceCacheIT {
+  private static final String DATABASE = "case_device_cache";
+  private static final String CASE_RESULT = "CASE WHEN src = 'match' THEN 'match' ELSE 'other' END";
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    // Keep every device in the same schema region and force a new batch for each device.
+    EnvFactory.getEnv().getConfig().getCommonConfig().setDefaultSchemaRegionGroupNumPerDatabase(1);
+    EnvFactory.getEnv().getConfig().getCommonConfig().setMaxTsBlockLineNumber(1);
+    EnvFactory.getEnv().initClusterEnvironment();
+    prepareTableData(new String[] {"CREATE DATABASE " + DATABASE});
+  }
+
+  @AfterClass
+  public static void tearDown() throws Exception {
+    EnvFactory.getEnv().cleanClusterEnvironment();
+  }
+
+  @Test
+  public void testDevicePredicatesAcrossBatches() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      createTable(statement, "predicate_cache");
+      String predicate = CASE_RESULT + " = src";
+      assertResultSetEqual(
+          statement.executeQuery("SHOW DEVICES FROM predicate_cache WHERE " + predicate),
+          "id,src,a,b,",
+          new HashSet<>(
+              Arrays.asList(
+                  "d0,match,null,null,",
+                  "d1,match,null,null,",
+                  "d2,other,null,null,",
+                  "d3,other,null,null,")));
+      assertResultSetEqual(
+          statement.executeQuery("COUNT DEVICES FROM predicate_cache WHERE " + predicate),
+          "count(devices),",
+          Collections.singleton("4,"));
+    }
+  }
+
+  @Test
+  public void testAttributeProjectionAcrossBatches() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      createTable(statement, "projection_cache");
+      statement.execute("UPDATE projection_cache SET a = " + CASE_RESULT + ", b = src");
+      assertResultSetEqual(
+          statement.executeQuery("SHOW DEVICES FROM projection_cache"),
+          "id,src,a,b,",
+          new HashSet<>(
+              Arrays.asList(
+                  "d0,match,match,match,",
+                  "d1,match,match,match,",
+                  "d2,other,other,other,",
+                  "d3,other,other,other,")));
+    }
+  }
+
+  @Test
+  public void testAttributeCommonExpressionClearedByCase() throws Exception {
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      createTable(statement, "common_cache");
+      statement.execute(
+          "UPDATE common_cache SET a = src "
+              + "WHERE CASE WHEN src = 'match' THEN true ELSE false END");
+      assertResultSetEqual(
+          statement.executeQuery("SHOW DEVICES FROM common_cache"),
+          "id,src,a,b,",
+          new HashSet<>(
+              Arrays.asList(
+                  "d0,match,match,null,",
+                  "d1,match,match,null,",
+                  "d2,other,null,null,",
+                  "d3,other,null,null,")));
+    }
+  }
+
+  private static void createTable(Statement statement, String table) throws SQLException {
+    statement.execute("USE " + DATABASE);
+    statement.execute(
+        "CREATE TABLE "
+            + table
+            + " (id STRING TAG, src STRING ATTRIBUTE, a STRING ATTRIBUTE, b STRING ATTRIBUTE, v DOUBLE FIELD)");
+    statement.execute(
+        "INSERT INTO "
+            + table
+            + "(time, id, src, v) VALUES (1, 'd0', 'match', 1), (2, 'd1', 'match', 1), "
+            + "(3, 'd2', 'other', 1), (4, 'd3', 'other', 1)");
+  }
+}

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/execution/operator/process/FilterAndProjectOperator.java

@@ -67,6 +67,10 @@ public class FilterAndProjectOperator implements ProcessOperator {
 
   private final List<ColumnTransformer> projectOutputTransformerList;
 
+  private final ColumnTransformer[] filterCacheTransformers;
+
+  private final ColumnTransformer[] projectCacheTransformers;
+
   private final TsBlockBuilder filterTsBlockBuilder;
 
   private final boolean hasNonMappableUDF;
@@ -100,6 +104,13 @@ public FilterAndProjectOperator(
     this.hasNonMappableUDF = hasNonMappableUDF;
     this.filterTsBlockBuilder = new TsBlockBuilder(8, filterOutputDataTypes);
     this.hasFilter = hasFilter;
+    List<ColumnTransformer> filterRoots = new ArrayList<>(commonTransformerList);
+    if (filterOutputTransformer != null) {
+      filterRoots.add(filterOutputTransformer);
+    }
+    this.filterCacheTransformers = ColumnTransformer.collectTransformers(filterRoots);
+    this.projectCacheTransformers =
+        ColumnTransformer.collectTransformers(projectOutputTransformerList);
   }
 
   public FilterAndProjectOperator(
@@ -111,6 +122,8 @@ public FilterAndProjectOperator(
     this.projectLeafColumnTransformerList =
         filterAndProjectOperator.projectLeafColumnTransformerList;
     this.projectOutputTransformerList = filterAndProjectOperator.projectOutputTransformerList;
+    this.filterCacheTransformers = filterAndProjectOperator.filterCacheTransformers;
+    this.projectCacheTransformers = filterAndProjectOperator.projectCacheTransformers;
     this.hasNonMappableUDF = filterAndProjectOperator.hasNonMappableUDF;
     this.hasFilter = filterAndProjectOperator.hasFilter;
     this.filterTsBlockBuilder = filterAndProjectOperator.filterTsBlockBuilder;
@@ -152,6 +165,12 @@ public TsBlock next() throws Exception {
   private TsBlock getFilterTsBlock(TsBlock input) {
     final Column originTimeColumn = input.getTimeColumn();
     final int positionCount = originTimeColumn.getPositionCount();
+    // Short-circuit evaluation can reset cache reference counts when a shared expression is
+    // recomputed. Invalidate the whole DAG before feeding a new block, including intermediate
+    // expressions whose caches may still have unconsumed references from the previous block.
+    for (ColumnTransformer columnTransformer : filterCacheTransformers) {
+      columnTransformer.clearCacheForNewInput();
+    }
     // feed Filter ColumnTransformer, including TimeStampColumnTransformer and constant
     for (LeafColumnTransformer leafColumnTransformer : filterLeafColumnTransformerList) {
       leafColumnTransformer.initFromTsBlock(input);
@@ -172,8 +191,8 @@ private TsBlock getFilterTsBlock(TsBlock input) {
     if (!hasNonMappableUDF) {
       // get result of calculated common sub expressions
       for (ColumnTransformer columnTransformer : commonTransformerList) {
-        // CASE WHEN clause would clear all its cache
-        // evaluate again to acquire cache
+        // CASE evaluation can clear a common expression after the filter used it. Recompute it
+        // before materializing the column for projection within this same input block.
         columnTransformer.tryEvaluate();
         resultColumns.add(columnTransformer.getColumn());
       }
@@ -263,6 +282,11 @@ public static boolean satisfy(Column filterColumn, int rowIndex) {
   private TsBlock getTransformedTsBlock(TsBlock input) {
     final Column originTimeColumn = input.getTimeColumn();
     final int positionCount = originTimeColumn.getPositionCount();
+    // The flat array was deduplicated at construction; clear each node before evaluating any
+    // output.
+    for (ColumnTransformer columnTransformer : projectCacheTransformers) {
+      columnTransformer.clearCacheForNewInput();
+    }
     // feed pre calculated data
     for (LeafColumnTransformer leafColumnTransformer : projectLeafColumnTransformerList) {
       leafColumnTransformer.initFromTsBlock(input);
@@ -439,6 +463,8 @@ public long ramBytesUsed() {
     return INSTANCE_SIZE
         + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(inputOperator)
         + MemoryEstimationHelper.getEstimatedSizeOfAccountableObject(operatorContext)
+        + RamUsageEstimator.shallowSizeOf(filterCacheTransformers)
+        + RamUsageEstimator.shallowSizeOf(projectCacheTransformers)
         + filterTsBlockBuilder.getRetainedSizeInBytes();
   }
 }

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/transformation/dag/column/AbstractCaseWhenThenColumnTransformer.java

@@ -30,6 +30,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Consumer;
 
 public abstract class AbstractCaseWhenThenColumnTransformer extends ColumnTransformer {
 
@@ -234,6 +235,22 @@ protected void checkType() {
     // do nothing
   }
 
+  @Override
+  protected boolean shouldClearChildrenCache() {
+    // Clearing a cached CASE during evaluation must preserve pending sibling columns.
+    return false;
+  }
+
+  @Override
+  protected void forEachChild(Consumer<ColumnTransformer> action) {
+    for (Pair<ColumnTransformer, ColumnTransformer> whenThenColumnTransformer :
+        whenThenTransformers) {
+      action.accept(whenThenColumnTransformer.left);
+      action.accept(whenThenColumnTransformer.right);
+    }
+    action.accept(elseTransformer);
+  }
+
   @Override
   public void close() {
     super.close();

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/transformation/dag/column/ColumnTransformer.java

@@ -25,6 +25,13 @@
 import org.apache.tsfile.read.common.type.Type;
 import org.apache.tsfile.read.common.type.TypeEnum;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.Set;
+import java.util.function.Consumer;
+
 import static org.apache.iotdb.calc.plan.relational.metadata.CommonMetadataUtils.isNumericType;
 
 public abstract class ColumnTransformer {
@@ -112,7 +119,54 @@ public void close() {
     // do nothing
   }
 
+  /**
+   * Clears caches during evaluation. Subclasses should define child edges via {@link
+   * #forEachChild(Consumer)} instead of overriding this method.
+   *
+   * <p>This legacy entry point remains non-final for compatibility with existing test doubles.
+   */
   public void clearCache() {
     this.columnCache.clear();
+    if (shouldClearChildrenCache()) {
+      forEachChild(ColumnTransformer::clearCache);
+    }
+  }
+
+  /**
+   * Clears only this node's cached column before initializing a new input block. Call this on every
+   * node collected by {@link #collectTransformers(Iterable)} before evaluating any expressions for
+   * the new block, when no sibling expression can still need a cached column.
+   */
+  public final void clearCacheForNewInput() {
+    this.columnCache.clear();
+  }
+
+  /**
+   * Collects all nodes reachable from the evaluation phase's roots, including CASE descendants,
+   * deduplicating by identity. Collect once after constructing the DAG and retain the flat array so
+   * new-input cache clearing needs no graph traversal or visited-set allocation.
+   */
+  public static ColumnTransformer[] collectTransformers(Iterable<ColumnTransformer> roots) {
+    List<ColumnTransformer> transformers = new ArrayList<>();
+    Set<ColumnTransformer> visited = Collections.newSetFromMap(new IdentityHashMap<>());
+    for (ColumnTransformer root : roots) {
+      root.collectTransformers(transformers, visited);
+    }
+    return transformers.toArray(new ColumnTransformer[0]);
+  }
+
+  private void collectTransformers(
+      List<ColumnTransformer> transformers, Set<ColumnTransformer> visited) {
+    if (visited.add(this)) {
+      transformers.add(this);
+      forEachChild(child -> child.collectTransformers(transformers, visited));
+    }
   }
+
+  protected boolean shouldClearChildrenCache() {
+    return true;
+  }
+
+  /** Defines the child edges used by both evaluation-time and new-input cache clearing. */
+  protected abstract void forEachChild(Consumer<ColumnTransformer> action);
 }

iotdb-core/calc-commons/src/main/java/org/apache/iotdb/calc/transformation/dag/column/FailFunctionColumnTransformer.java

@@ -24,6 +24,8 @@
 import org.apache.tsfile.read.common.block.column.NullColumn;
 import org.apache.tsfile.read.common.type.Type;
 
+import java.util.function.Consumer;
+
 /**
  * when evaluate of FailFunctionColumnTransformer is called, throw Exception to inform the user with
  * specified errorMsg
@@ -59,4 +61,9 @@ public void evaluateWithSelection(boolean[] selection) {
   protected void checkType() {
     // do nothing
   }
+
+  @Override
+  protected void forEachChild(Consumer<ColumnTransformer> action) {
+    // The fail function has no children.
+  }
 }