github-advanced-security[bot] commented on code in PR #19729: URL: https://github.com/apache/druid/pull/19729#discussion_r3635585251
########## sql/src/main/java/org/apache/druid/sql/calcite/rule/SegmentsRollupRule.java: ########## @@ -0,0 +1,600 @@ +/* + * 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.druid.sql.calcite.rule; + +import com.google.common.collect.BoundType; +import com.google.common.collect.Range; +import org.apache.calcite.interpreter.BindableConvention; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.core.Project; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.util.Sarg; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.rel.SegmentsRollupRel; +import org.apache.druid.sql.calcite.schema.DatasourceSegmentStats; +import org.apache.druid.sql.calcite.schema.SegmentsRollup; +import org.apache.druid.sql.calcite.schema.SystemSchema; +import org.apache.druid.sql.calcite.schema.SystemSchema.SegmentsRollupSource; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.function.Function; + +/** + * Rewrites a per-datasource aggregate over {@code sys.segments} - + * {@code SELECT datasource, <aggregates> ... GROUP BY datasource} - into a {@link SegmentsRollupRel} + * that reads the precomputed {@link SegmentsRollup} status sub-cube, so the query is O(#datasources) + * instead of a full segment scan. + * + * <p>It recognizes any aggregate that is: + * <ul> + * <li>{@code COUNT(*)}, {@code SUM}/{@code AVG}({@code size} | {@code num_rows} | + * {@code size*num_replicas}), or {@code MIN}/{@code MAX}({@code num_rows}); and</li> + * <li>optionally {@code FILTER}ed by an arbitrary boolean expression over the low-cardinality + * <em>status</em> columns ({@code is_active, is_published, is_available, is_realtime, + * is_overshadowed}, and {@code replication_factor} compared against 0). The filter is turned + * into a mask of cube cells by evaluating it over the finite status domain.</li> + * </ul> + * + * <p>Anything it does not fully recognize - a different grouping, a filter touching a non-status + * column, an unsupported measure/operator, or a rollup that is not yet populated - leaves the tree + * unchanged, so the query falls back to the normal scan and stays correct. It runs as a Hep pre-pass + * (before {@code AVG} is reduced to {@code SUM/COUNT}), so {@code AVG} is still a single call here. + */ +public class SegmentsRollupRule extends RelOptRule +{ + // Column positions in SystemSchema.SEGMENTS_SIGNATURE, derived from the signature (not hardcoded) so + // they track the column order. Valid because the rule only fires on a table marked SegmentsRollupSource. + private static final int COL_DATASOURCE = SegmentsRollupSource.COL_DATASOURCE; + private static final int COL_SIZE = SegmentsRollupSource.COL_SIZE; + private static final int COL_NUM_REPLICAS = SegmentsRollupSource.COL_NUM_REPLICAS; + private static final int COL_NUM_ROWS = SegmentsRollupSource.COL_NUM_ROWS; + private static final int COL_IS_ACTIVE = SegmentsRollupSource.COL_IS_ACTIVE; + private static final int COL_IS_PUBLISHED = SegmentsRollupSource.COL_IS_PUBLISHED; + private static final int COL_IS_AVAILABLE = SegmentsRollupSource.COL_IS_AVAILABLE; + private static final int COL_IS_REALTIME = SegmentsRollupSource.COL_IS_REALTIME; + private static final int COL_IS_OVERSHADOWED = SegmentsRollupSource.COL_IS_OVERSHADOWED; + private static final int COL_REPLICATION_FACTOR = SegmentsRollupSource.COL_REPLICATION_FACTOR; + + /** Mask selecting every cell - used for an aggregate with no FILTER. */ + private static final boolean[] ALL_CELLS = allCells(); + + private final SegmentsRollup rollup; + private final AuthorizerMapper authorizerMapper; + private final PlannerContext plannerContext; + + public SegmentsRollupRule( + final SegmentsRollup rollup, + final AuthorizerMapper authorizerMapper, + final PlannerContext plannerContext + ) + { + super( + operand(Aggregate.class, operand(Project.class, operand(TableScan.class, none()))), Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [RelOptRule.operand](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/11356) ########## sql/src/main/java/org/apache/druid/sql/calcite/rule/SegmentsRollupRule.java: ########## @@ -0,0 +1,600 @@ +/* + * 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.druid.sql.calcite.rule; + +import com.google.common.collect.BoundType; +import com.google.common.collect.Range; +import org.apache.calcite.interpreter.BindableConvention; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.core.Project; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.util.Sarg; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.rel.SegmentsRollupRel; +import org.apache.druid.sql.calcite.schema.DatasourceSegmentStats; +import org.apache.druid.sql.calcite.schema.SegmentsRollup; +import org.apache.druid.sql.calcite.schema.SystemSchema; +import org.apache.druid.sql.calcite.schema.SystemSchema.SegmentsRollupSource; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.function.Function; + +/** + * Rewrites a per-datasource aggregate over {@code sys.segments} - + * {@code SELECT datasource, <aggregates> ... GROUP BY datasource} - into a {@link SegmentsRollupRel} + * that reads the precomputed {@link SegmentsRollup} status sub-cube, so the query is O(#datasources) + * instead of a full segment scan. + * + * <p>It recognizes any aggregate that is: + * <ul> + * <li>{@code COUNT(*)}, {@code SUM}/{@code AVG}({@code size} | {@code num_rows} | + * {@code size*num_replicas}), or {@code MIN}/{@code MAX}({@code num_rows}); and</li> + * <li>optionally {@code FILTER}ed by an arbitrary boolean expression over the low-cardinality + * <em>status</em> columns ({@code is_active, is_published, is_available, is_realtime, + * is_overshadowed}, and {@code replication_factor} compared against 0). The filter is turned + * into a mask of cube cells by evaluating it over the finite status domain.</li> + * </ul> + * + * <p>Anything it does not fully recognize - a different grouping, a filter touching a non-status + * column, an unsupported measure/operator, or a rollup that is not yet populated - leaves the tree + * unchanged, so the query falls back to the normal scan and stays correct. It runs as a Hep pre-pass + * (before {@code AVG} is reduced to {@code SUM/COUNT}), so {@code AVG} is still a single call here. + */ +public class SegmentsRollupRule extends RelOptRule +{ + // Column positions in SystemSchema.SEGMENTS_SIGNATURE, derived from the signature (not hardcoded) so + // they track the column order. Valid because the rule only fires on a table marked SegmentsRollupSource. + private static final int COL_DATASOURCE = SegmentsRollupSource.COL_DATASOURCE; + private static final int COL_SIZE = SegmentsRollupSource.COL_SIZE; + private static final int COL_NUM_REPLICAS = SegmentsRollupSource.COL_NUM_REPLICAS; + private static final int COL_NUM_ROWS = SegmentsRollupSource.COL_NUM_ROWS; + private static final int COL_IS_ACTIVE = SegmentsRollupSource.COL_IS_ACTIVE; + private static final int COL_IS_PUBLISHED = SegmentsRollupSource.COL_IS_PUBLISHED; + private static final int COL_IS_AVAILABLE = SegmentsRollupSource.COL_IS_AVAILABLE; + private static final int COL_IS_REALTIME = SegmentsRollupSource.COL_IS_REALTIME; + private static final int COL_IS_OVERSHADOWED = SegmentsRollupSource.COL_IS_OVERSHADOWED; + private static final int COL_REPLICATION_FACTOR = SegmentsRollupSource.COL_REPLICATION_FACTOR; + + /** Mask selecting every cell - used for an aggregate with no FILTER. */ + private static final boolean[] ALL_CELLS = allCells(); + + private final SegmentsRollup rollup; + private final AuthorizerMapper authorizerMapper; + private final PlannerContext plannerContext; + + public SegmentsRollupRule( + final SegmentsRollup rollup, + final AuthorizerMapper authorizerMapper, + final PlannerContext plannerContext + ) + { + super( + operand(Aggregate.class, operand(Project.class, operand(TableScan.class, none()))), Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [RelOptRule.operand](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/11358) ########## sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java: ########## @@ -307,11 +315,149 @@ return tableMap; } + /** + * Per-segment status/measure fields derived identically for {@code sys.segments} rows and the + * internal per-datasource rollup (see {@code SegmentsRollup}), so the two can never drift. Encoding + * of the boolean flags into the row's LONG columns is left to the caller. {@link #replicationFactor} + * is already encoded ({@link #REPLICATION_FACTOR_UNKNOWN} when unknown), matching the row column, so + * counters that key off {@code replication_factor} agree with a query over sys.segments. + */ + static final class DerivedSegmentStatus + { + final long numReplicas; + final long numRows; + final boolean isAvailable; + final boolean isRealtime; + final boolean isPublished; + final boolean isActive; + final boolean isOvershadowed; + final long replicationFactor; + + DerivedSegmentStatus( + long numReplicas, + long numRows, + boolean isAvailable, + boolean isRealtime, + boolean isPublished, + boolean isActive, + boolean isOvershadowed, + long replicationFactor + ) + { + this.numReplicas = numReplicas; + this.numRows = numRows; + this.isAvailable = isAvailable; + this.isRealtime = isRealtime; + this.isPublished = isPublished; + this.isActive = isActive; + this.isOvershadowed = isOvershadowed; + this.replicationFactor = replicationFactor; + } + + /** + * Derivation for a segment from the Coordinator's published (+ realtime, with centralized schema) + * set, joined with any available metadata the broker has for it. + */ + static DerivedSegmentStatus forPublished( + final SegmentStatusInCluster val, + @Nullable final AvailableSegmentMetadata available + ) + { + final DataSegment segment = val.getDataSegment(); + + long numReplicas = 0L; + boolean isAvailable = false; + if (available != null) { + numReplicas = available.getNumReplicas(); + isAvailable = available.getNumReplicas() > 0; + } + + final long numRows; + if (segment.getTotalRows() != null) { + // the recent version of DataSegment stores numRows + numRows = segment.getTotalRows(); + } else if (val.getNumRows() != null) { + // If druid.centralizedDatasourceSchema.enabled is set on the Coordinator, SegmentMetadataCache + // on the broker might have outdated or no information regarding numRows and rowSignature for a + // segment. In that case, we use numRows from the segment polled from the Coordinator. + numRows = val.getNumRows(); + } else if (available != null) { + numRows = available.getNumRows(); + } else { + numRows = 0L; + } + + final boolean isRealtime = val.isRealtime(); + // A segment from this set is published unless it is a realtime segment (mutually exclusive). + final boolean isPublished = !val.isRealtime(); + // is_active is true for published segments that are not overshadowed, or else realtime segments. + final boolean isActive = isPublished ? !val.isOvershadowed() : val.isRealtime(); + final long replicationFactor = + val.getReplicationFactor() == null ? REPLICATION_FACTOR_UNKNOWN : val.getReplicationFactor(); + + return new DerivedSegmentStatus( + numReplicas, + numRows, + isAvailable, + isRealtime, + isPublished, + isActive, + val.isOvershadowed(), + replicationFactor + ); + } + + /** + * Derivation for a segment known only from the broker's available (served) view - not present in + * the published set. Assumed available, unpublished, non-overshadowed. + */ + static DerivedSegmentStatus forAvailable(final AvailableSegmentMetadata val) + { + final DataSegment segment = val.getSegment(); + final long numRows = segment.getTotalRows() != null ? segment.getTotalRows() : val.getNumRows(); + // AvailableSegmentMetadata.isRealtime() returns a long flag (0/1), unlike SegmentStatusInCluster. + final boolean isRealtime = val.isRealtime() != 0; + return new DerivedSegmentStatus( + val.getNumReplicas(), + numRows, + true, + isRealtime, + false, + // is_active is true for unpublished segments iff they are realtime. + isRealtime, + false, + REPLICATION_FACTOR_UNKNOWN + ); + } + } + + /** + * Marker for the {@code sys.segments} table so an optimizer rule can recognize a scan of it, plus + * the {@link #SEGMENTS_SIGNATURE} column positions the rule needs - derived from the signature here + * (rather than hardcoded in the rule) so they can never silently drift if the column order changes. + * The rule depends on this interface instead of the package-private {@link SegmentsTable}. + */ + public interface SegmentsRollupSource + { + int COL_DATASOURCE = SEGMENTS_SIGNATURE.indexOf("datasource"); + int COL_SIZE = SEGMENTS_SIGNATURE.indexOf("size"); + int COL_NUM_REPLICAS = SEGMENTS_SIGNATURE.indexOf("num_replicas"); + int COL_NUM_ROWS = SEGMENTS_SIGNATURE.indexOf("num_rows"); + int COL_IS_ACTIVE = SEGMENTS_SIGNATURE.indexOf("is_active"); + int COL_IS_PUBLISHED = SEGMENTS_SIGNATURE.indexOf("is_published"); + int COL_IS_AVAILABLE = SEGMENTS_SIGNATURE.indexOf("is_available"); + int COL_IS_REALTIME = SEGMENTS_SIGNATURE.indexOf("is_realtime"); + int COL_IS_OVERSHADOWED = SEGMENTS_SIGNATURE.indexOf("is_overshadowed"); + int COL_REPLICATION_FACTOR = SEGMENTS_SIGNATURE.indexOf("replication_factor"); + } + /** * This table contains row per segment from metadata store as well as served segments. */ - static class SegmentsTable extends AbstractTable implements ProjectableFilterableTable + static class SegmentsTable extends AbstractTable implements ProjectableFilterableTable, SegmentsRollupSource Review Comment: ## CodeQL / Constant interface anti-pattern Type SegmentsTable implements constant interface [SegmentsRollupSource](1). [Show more details](https://github.com/apache/druid/security/code-scanning/11355) ########## sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java: ########## @@ -307,11 +315,149 @@ return tableMap; } + /** + * Per-segment status/measure fields derived identically for {@code sys.segments} rows and the + * internal per-datasource rollup (see {@code SegmentsRollup}), so the two can never drift. Encoding + * of the boolean flags into the row's LONG columns is left to the caller. {@link #replicationFactor} + * is already encoded ({@link #REPLICATION_FACTOR_UNKNOWN} when unknown), matching the row column, so + * counters that key off {@code replication_factor} agree with a query over sys.segments. + */ + static final class DerivedSegmentStatus + { + final long numReplicas; + final long numRows; + final boolean isAvailable; + final boolean isRealtime; + final boolean isPublished; + final boolean isActive; + final boolean isOvershadowed; + final long replicationFactor; + + DerivedSegmentStatus( + long numReplicas, + long numRows, + boolean isAvailable, + boolean isRealtime, + boolean isPublished, + boolean isActive, + boolean isOvershadowed, + long replicationFactor + ) + { + this.numReplicas = numReplicas; + this.numRows = numRows; + this.isAvailable = isAvailable; + this.isRealtime = isRealtime; + this.isPublished = isPublished; + this.isActive = isActive; + this.isOvershadowed = isOvershadowed; + this.replicationFactor = replicationFactor; + } + + /** + * Derivation for a segment from the Coordinator's published (+ realtime, with centralized schema) + * set, joined with any available metadata the broker has for it. + */ + static DerivedSegmentStatus forPublished( + final SegmentStatusInCluster val, + @Nullable final AvailableSegmentMetadata available + ) + { + final DataSegment segment = val.getDataSegment(); + + long numReplicas = 0L; + boolean isAvailable = false; + if (available != null) { + numReplicas = available.getNumReplicas(); + isAvailable = available.getNumReplicas() > 0; + } + + final long numRows; + if (segment.getTotalRows() != null) { + // the recent version of DataSegment stores numRows + numRows = segment.getTotalRows(); + } else if (val.getNumRows() != null) { + // If druid.centralizedDatasourceSchema.enabled is set on the Coordinator, SegmentMetadataCache + // on the broker might have outdated or no information regarding numRows and rowSignature for a + // segment. In that case, we use numRows from the segment polled from the Coordinator. + numRows = val.getNumRows(); + } else if (available != null) { + numRows = available.getNumRows(); + } else { + numRows = 0L; + } + + final boolean isRealtime = val.isRealtime(); + // A segment from this set is published unless it is a realtime segment (mutually exclusive). + final boolean isPublished = !val.isRealtime(); + // is_active is true for published segments that are not overshadowed, or else realtime segments. + final boolean isActive = isPublished ? !val.isOvershadowed() : val.isRealtime(); + final long replicationFactor = + val.getReplicationFactor() == null ? REPLICATION_FACTOR_UNKNOWN : val.getReplicationFactor(); + + return new DerivedSegmentStatus( + numReplicas, + numRows, + isAvailable, + isRealtime, + isPublished, + isActive, + val.isOvershadowed(), + replicationFactor + ); + } + + /** + * Derivation for a segment known only from the broker's available (served) view - not present in + * the published set. Assumed available, unpublished, non-overshadowed. + */ + static DerivedSegmentStatus forAvailable(final AvailableSegmentMetadata val) + { + final DataSegment segment = val.getSegment(); + final long numRows = segment.getTotalRows() != null ? segment.getTotalRows() : val.getNumRows(); Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [AvailableSegmentMetadata.getNumRows](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/11361) ########## sql/src/main/java/org/apache/druid/sql/calcite/rel/SegmentsRollupRel.java: ########## @@ -0,0 +1,165 @@ +/* + * 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.druid.sql.calcite.rel; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSortedMap; +import org.apache.calcite.DataContext; +import org.apache.calcite.interpreter.BindableRel; +import org.apache.calcite.interpreter.InterpretableRel; +import org.apache.calcite.interpreter.Interpreter; +import org.apache.calcite.interpreter.Node; +import org.apache.calcite.interpreter.Row; +import org.apache.calcite.interpreter.Sink; +import org.apache.calcite.linq4j.Enumerable; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptCost; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.RelTraitSet; +import org.apache.calcite.rel.AbstractRelNode; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.metadata.RelMetadataQuery; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizationUtils; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.security.ResourceAction; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.schema.DatasourceSegmentStats; +import org.apache.druid.sql.calcite.schema.SegmentsRollup; + +import java.util.Collections; +import java.util.List; +import java.util.function.Function; + +/** + * A leaf {@link BindableRel} that answers the console's {@code GROUP BY datasource} aggregate over + * {@code sys.segments} from the precomputed {@link SegmentsRollup}, instead of scanning every segment. + * Produced by {@code SegmentsRollupRule}. Each row is {@code [datasource, agg0, agg1, ...]}: the group + * key followed by one value per aggregate call, in the order Calcite lays out the {@code Aggregate} + * output. Per-datasource read authorization is applied here so results match a scan of + * {@code sys.segments}. + */ +public class SegmentsRollupRel extends AbstractRelNode implements BindableRel Review Comment: ## CodeQL / No clone method No clone method, yet implements Cloneable. [Show more details](https://github.com/apache/druid/security/code-scanning/11362) ########## sql/src/main/java/org/apache/druid/sql/calcite/rule/SegmentsRollupRule.java: ########## @@ -0,0 +1,600 @@ +/* + * 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.druid.sql.calcite.rule; + +import com.google.common.collect.BoundType; +import com.google.common.collect.Range; +import org.apache.calcite.interpreter.BindableConvention; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.core.Project; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.util.Sarg; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.rel.SegmentsRollupRel; +import org.apache.druid.sql.calcite.schema.DatasourceSegmentStats; +import org.apache.druid.sql.calcite.schema.SegmentsRollup; +import org.apache.druid.sql.calcite.schema.SystemSchema; +import org.apache.druid.sql.calcite.schema.SystemSchema.SegmentsRollupSource; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.function.Function; + +/** + * Rewrites a per-datasource aggregate over {@code sys.segments} - + * {@code SELECT datasource, <aggregates> ... GROUP BY datasource} - into a {@link SegmentsRollupRel} + * that reads the precomputed {@link SegmentsRollup} status sub-cube, so the query is O(#datasources) + * instead of a full segment scan. + * + * <p>It recognizes any aggregate that is: + * <ul> + * <li>{@code COUNT(*)}, {@code SUM}/{@code AVG}({@code size} | {@code num_rows} | + * {@code size*num_replicas}), or {@code MIN}/{@code MAX}({@code num_rows}); and</li> + * <li>optionally {@code FILTER}ed by an arbitrary boolean expression over the low-cardinality + * <em>status</em> columns ({@code is_active, is_published, is_available, is_realtime, + * is_overshadowed}, and {@code replication_factor} compared against 0). The filter is turned + * into a mask of cube cells by evaluating it over the finite status domain.</li> + * </ul> + * + * <p>Anything it does not fully recognize - a different grouping, a filter touching a non-status + * column, an unsupported measure/operator, or a rollup that is not yet populated - leaves the tree + * unchanged, so the query falls back to the normal scan and stays correct. It runs as a Hep pre-pass + * (before {@code AVG} is reduced to {@code SUM/COUNT}), so {@code AVG} is still a single call here. + */ +public class SegmentsRollupRule extends RelOptRule +{ + // Column positions in SystemSchema.SEGMENTS_SIGNATURE, derived from the signature (not hardcoded) so + // they track the column order. Valid because the rule only fires on a table marked SegmentsRollupSource. + private static final int COL_DATASOURCE = SegmentsRollupSource.COL_DATASOURCE; + private static final int COL_SIZE = SegmentsRollupSource.COL_SIZE; + private static final int COL_NUM_REPLICAS = SegmentsRollupSource.COL_NUM_REPLICAS; + private static final int COL_NUM_ROWS = SegmentsRollupSource.COL_NUM_ROWS; + private static final int COL_IS_ACTIVE = SegmentsRollupSource.COL_IS_ACTIVE; + private static final int COL_IS_PUBLISHED = SegmentsRollupSource.COL_IS_PUBLISHED; + private static final int COL_IS_AVAILABLE = SegmentsRollupSource.COL_IS_AVAILABLE; + private static final int COL_IS_REALTIME = SegmentsRollupSource.COL_IS_REALTIME; + private static final int COL_IS_OVERSHADOWED = SegmentsRollupSource.COL_IS_OVERSHADOWED; + private static final int COL_REPLICATION_FACTOR = SegmentsRollupSource.COL_REPLICATION_FACTOR; + + /** Mask selecting every cell - used for an aggregate with no FILTER. */ + private static final boolean[] ALL_CELLS = allCells(); + + private final SegmentsRollup rollup; + private final AuthorizerMapper authorizerMapper; + private final PlannerContext plannerContext; + + public SegmentsRollupRule( + final SegmentsRollup rollup, + final AuthorizerMapper authorizerMapper, + final PlannerContext plannerContext + ) + { + super( + operand(Aggregate.class, operand(Project.class, operand(TableScan.class, none()))), Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [RelOptRule.none](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/11359) ########## sql/src/main/java/org/apache/druid/sql/calcite/schema/SystemSchema.java: ########## @@ -307,11 +315,149 @@ return tableMap; } + /** + * Per-segment status/measure fields derived identically for {@code sys.segments} rows and the + * internal per-datasource rollup (see {@code SegmentsRollup}), so the two can never drift. Encoding + * of the boolean flags into the row's LONG columns is left to the caller. {@link #replicationFactor} + * is already encoded ({@link #REPLICATION_FACTOR_UNKNOWN} when unknown), matching the row column, so + * counters that key off {@code replication_factor} agree with a query over sys.segments. + */ + static final class DerivedSegmentStatus + { + final long numReplicas; + final long numRows; + final boolean isAvailable; + final boolean isRealtime; + final boolean isPublished; + final boolean isActive; + final boolean isOvershadowed; + final long replicationFactor; + + DerivedSegmentStatus( + long numReplicas, + long numRows, + boolean isAvailable, + boolean isRealtime, + boolean isPublished, + boolean isActive, + boolean isOvershadowed, + long replicationFactor + ) + { + this.numReplicas = numReplicas; + this.numRows = numRows; + this.isAvailable = isAvailable; + this.isRealtime = isRealtime; + this.isPublished = isPublished; + this.isActive = isActive; + this.isOvershadowed = isOvershadowed; + this.replicationFactor = replicationFactor; + } + + /** + * Derivation for a segment from the Coordinator's published (+ realtime, with centralized schema) + * set, joined with any available metadata the broker has for it. + */ + static DerivedSegmentStatus forPublished( + final SegmentStatusInCluster val, + @Nullable final AvailableSegmentMetadata available + ) + { + final DataSegment segment = val.getDataSegment(); + + long numReplicas = 0L; + boolean isAvailable = false; + if (available != null) { + numReplicas = available.getNumReplicas(); + isAvailable = available.getNumReplicas() > 0; + } + + final long numRows; + if (segment.getTotalRows() != null) { + // the recent version of DataSegment stores numRows + numRows = segment.getTotalRows(); + } else if (val.getNumRows() != null) { + // If druid.centralizedDatasourceSchema.enabled is set on the Coordinator, SegmentMetadataCache + // on the broker might have outdated or no information regarding numRows and rowSignature for a + // segment. In that case, we use numRows from the segment polled from the Coordinator. + numRows = val.getNumRows(); + } else if (available != null) { + numRows = available.getNumRows(); Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [AvailableSegmentMetadata.getNumRows](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/11360) ########## sql/src/main/java/org/apache/druid/sql/calcite/rule/SegmentsRollupRule.java: ########## @@ -0,0 +1,600 @@ +/* + * 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.druid.sql.calcite.rule; + +import com.google.common.collect.BoundType; +import com.google.common.collect.Range; +import org.apache.calcite.interpreter.BindableConvention; +import org.apache.calcite.plan.RelOptRule; +import org.apache.calcite.plan.RelOptRuleCall; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.core.Project; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.util.Sarg; +import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.sql.calcite.planner.PlannerContext; +import org.apache.druid.sql.calcite.rel.SegmentsRollupRel; +import org.apache.druid.sql.calcite.schema.DatasourceSegmentStats; +import org.apache.druid.sql.calcite.schema.SegmentsRollup; +import org.apache.druid.sql.calcite.schema.SystemSchema; +import org.apache.druid.sql.calcite.schema.SystemSchema.SegmentsRollupSource; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.function.Function; + +/** + * Rewrites a per-datasource aggregate over {@code sys.segments} - + * {@code SELECT datasource, <aggregates> ... GROUP BY datasource} - into a {@link SegmentsRollupRel} + * that reads the precomputed {@link SegmentsRollup} status sub-cube, so the query is O(#datasources) + * instead of a full segment scan. + * + * <p>It recognizes any aggregate that is: + * <ul> + * <li>{@code COUNT(*)}, {@code SUM}/{@code AVG}({@code size} | {@code num_rows} | + * {@code size*num_replicas}), or {@code MIN}/{@code MAX}({@code num_rows}); and</li> + * <li>optionally {@code FILTER}ed by an arbitrary boolean expression over the low-cardinality + * <em>status</em> columns ({@code is_active, is_published, is_available, is_realtime, + * is_overshadowed}, and {@code replication_factor} compared against 0). The filter is turned + * into a mask of cube cells by evaluating it over the finite status domain.</li> + * </ul> + * + * <p>Anything it does not fully recognize - a different grouping, a filter touching a non-status + * column, an unsupported measure/operator, or a rollup that is not yet populated - leaves the tree + * unchanged, so the query falls back to the normal scan and stays correct. It runs as a Hep pre-pass + * (before {@code AVG} is reduced to {@code SUM/COUNT}), so {@code AVG} is still a single call here. + */ +public class SegmentsRollupRule extends RelOptRule +{ + // Column positions in SystemSchema.SEGMENTS_SIGNATURE, derived from the signature (not hardcoded) so + // they track the column order. Valid because the rule only fires on a table marked SegmentsRollupSource. + private static final int COL_DATASOURCE = SegmentsRollupSource.COL_DATASOURCE; + private static final int COL_SIZE = SegmentsRollupSource.COL_SIZE; + private static final int COL_NUM_REPLICAS = SegmentsRollupSource.COL_NUM_REPLICAS; + private static final int COL_NUM_ROWS = SegmentsRollupSource.COL_NUM_ROWS; + private static final int COL_IS_ACTIVE = SegmentsRollupSource.COL_IS_ACTIVE; + private static final int COL_IS_PUBLISHED = SegmentsRollupSource.COL_IS_PUBLISHED; + private static final int COL_IS_AVAILABLE = SegmentsRollupSource.COL_IS_AVAILABLE; + private static final int COL_IS_REALTIME = SegmentsRollupSource.COL_IS_REALTIME; + private static final int COL_IS_OVERSHADOWED = SegmentsRollupSource.COL_IS_OVERSHADOWED; + private static final int COL_REPLICATION_FACTOR = SegmentsRollupSource.COL_REPLICATION_FACTOR; + + /** Mask selecting every cell - used for an aggregate with no FILTER. */ + private static final boolean[] ALL_CELLS = allCells(); + + private final SegmentsRollup rollup; + private final AuthorizerMapper authorizerMapper; + private final PlannerContext plannerContext; + + public SegmentsRollupRule( + final SegmentsRollup rollup, + final AuthorizerMapper authorizerMapper, + final PlannerContext plannerContext + ) + { + super( + operand(Aggregate.class, operand(Project.class, operand(TableScan.class, none()))), Review Comment: ## CodeQL / Deprecated method or constructor invocation Invoking [RelOptRule.operand](1) should be avoided because it has been deprecated. [Show more details](https://github.com/apache/druid/security/code-scanning/11357) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
