FrankChen021 commented on code in PR #19379: URL: https://github.com/apache/druid/pull/19379#discussion_r4071755616
########## processing/src/main/java/org/apache/druid/segment/transform/ScanTransformer.java: ########## @@ -0,0 +1,339 @@ +/* + * 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.segment.transform; + +import org.apache.druid.data.input.InputRow; +import org.apache.druid.data.input.InputRowListPlusRawValues; +import org.apache.druid.data.input.ListBasedInputRow; +import org.apache.druid.data.input.MapBasedInputRow; +import org.apache.druid.error.DruidException; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.query.DataSource; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.query.UnnestDataSource; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.segment.ColumnSelectorFactory; +import org.apache.druid.segment.ColumnValueSelector; +import org.apache.druid.segment.Cursor; +import org.apache.druid.segment.CursorBuildSpec; +import org.apache.druid.segment.CursorFactory; +import org.apache.druid.segment.CursorHolder; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.SegmentMapFunction; +import org.apache.druid.segment.VirtualColumn; +import org.apache.druid.segment.column.ColumnHolder; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.filter.Filters; +import org.apache.druid.timeline.SegmentId; +import org.joda.time.Interval; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * A {@link BaseTransformer} that processes input rows through a reusable scan query cursor pipeline. + * + * <p>The pipeline is built once at construction: a {@link SettableRowCursorFactory} is wrapped by the + * scan query's {@link SegmentMapFunction} (e.g., unnest, filter). For each input row, the row is set + * on the factory and the cursor is {@link Cursor#reset reset} — no per-row segment or cursor allocation. + * + * <p>When the scan query produces zero output rows (e.g., null/missing arrays, or filter rejection), + * the input row is dropped. This matches native Druid UNNEST / CROSS JOIN semantics where + * null or empty arrays produce zero rows. + * + * <p>This class is not thread-safe. Each reader thread should have its own instance. + */ +public class ScanTransformer implements BaseTransformer +{ + private final ScanQuery query; + // Field names that must never be promoted to a dimension, regardless of whether they show up in this + // scan query's result columns. See resolveDimensionColumns() for why this can't be derived locally. + private final Set<String> dimensionExclusions; + private final SettableRowCursorFactory baseCursorFactory; + private final CursorHolder cursorHolder; + private Cursor cursor; + + ScanTransformer(final ScanQuery scanQuery, final Set<String> dimensionExclusions) + { + this.query = scanQuery.withOverriddenContext( + Map.of(QueryContexts.TIMEOUT_KEY, 0) + ); + this.dimensionExclusions = dimensionExclusions; + + final RowSignature broadSignature = RowSignature.builder() + .add(ColumnHolder.TIME_COLUMN_NAME, ColumnType.LONG) + .build(); + + final CursorBuildSpec cursorBuildSpec = CursorBuildSpec.builder() + .setInterval(query.getSingleInterval()) + .setFilter(Filters.toFilter(query.getFilter())) + .setVirtualColumns(query.getVirtualColumns()) + .build(); + + this.baseCursorFactory = new SettableRowCursorFactory(broadSignature); + final SegmentMapFunction segmentMapFunction = query.getDataSource().createSegmentMapFunction(query); + final Segment mappedSegment = segmentMapFunction.apply(Optional.of(new CursorFactorySegment(baseCursorFactory))) + .orElseThrow(() -> new ISE("SegmentMapFunction returned empty")); + final CursorFactory mappedCursorFactory = mappedSegment.as(CursorFactory.class); + this.cursorHolder = mappedCursorFactory.makeCursorHolder(cursorBuildSpec); + } + + @Override + public boolean hasMultiRowTransform() + { + return true; + } + + @Override + @Nullable + public InputRow transform(@Nullable final InputRow row) + { + throw new UnsupportedOperationException( + "ScanTransformer does not support single-row transform; use transformToList()" + ); + } + + @Override + public List<InputRow> transformToList(@Nullable final InputRow row) + { + if (row == null) { + return List.of(); + } + + return process(row); + } + + @Override + @Nullable + public InputRowListPlusRawValues transform(@Nullable final InputRowListPlusRawValues row) + { + if (row == null || row.getInputRows() == null) { + return row; + } + + final List<InputRow> inputRows = row.getInputRows(); + final List<Map<String, Object>> inputRawValues = row.getRawValuesList(); + final List<InputRow> outputRows = new ArrayList<>(); + final List<Map<String, Object>> outputRawValues = inputRawValues == null ? null : new ArrayList<>(); + + for (int i = 0; i < inputRows.size(); i++) { + final List<InputRow> expandedRows = transformToList(inputRows.get(i)); + outputRows.addAll(expandedRows); + if (outputRawValues != null) { + for (int j = 0; j < expandedRows.size(); j++) { + outputRawValues.add(inputRawValues.get(i)); + } + } + } + + return InputRowListPlusRawValues.ofList(outputRawValues, outputRows, row.getParseException()); + } + + @Override + public void close() throws IOException + { + cursorHolder.close(); + } + + private List<InputRow> process(final InputRow inputRow) + { + baseCursorFactory.set(inputRow); + + if (cursor == null) { + cursor = cursorHolder.asCursor(); + } else { + cursor.reset(); + } + + if (cursor == null || cursor.isDone()) { + return List.of(); + } + + final Set<String> nonDimensionEventFields = resolveNonDimensionEventFields(inputRow); + final List<String> columns = resolveColumnsForRow(inputRow, nonDimensionEventFields); + final List<String> dimensionColumns = resolveDimensionColumns(inputRow, columns, nonDimensionEventFields); + final ColumnSelectorFactory selectorFactory = cursor.getColumnSelectorFactory(); + + // Selectors are lazy views over the cursor's current position — create them once per column + // here, then re-read via getObject() as the cursor advances, rather than reallocating a selector + // for every (output-row x column) pair. + final ColumnValueSelector<?>[] selectors = new ColumnValueSelector<?>[columns.size()]; + for (int i = 0; i < columns.size(); i++) { + selectors[i] = selectorFactory.makeColumnValueSelector(columns.get(i)); + } + + // The query re-executes fresh for each input row (cursor reset above), so this row's own unnested + // expansion is the query's entire result set for this execution — offset/limit bound that set, + // the same way they'd bound any other scan query's result set. This is necessarily per input row + // rather than global across the ingestion job: there is no single ordered stream spanning rows + // (let alone across the parallel/rolling readers of a real ingestion job) for them to paginate. + final long offset = query.getScanRowsOffset(); + final long limit = query.getScanRowsLimit(); + long skipped = 0; + long emitted = 0; + + final List<InputRow> result = new ArrayList<>(); + while (!cursor.isDone() && emitted < limit) { + if (skipped < offset) { + skipped++; + cursor.advance(); + continue; + } + final Map<String, Object> event = new LinkedHashMap<>(); + for (int i = 0; i < columns.size(); i++) { + event.put(columns.get(i), selectors[i].getObject()); + } + result.add(new MapBasedInputRow(inputRow.getTimestampFromEpoch(), dimensionColumns, event)); + emitted++; + cursor.advance(); + } + + return result; + } + + /** + * Returns the raw fields present on {@code inputRow} that are not in {@link InputRow#getDimensions()} + * — e.g. metric inputs that {@code DataSchema} added to dimensionExclusions. These must still be read + * into the expanded rows' event maps (for aggregators), but must not be promoted to dimensions. + * + * <p>Handles both {@link MapBasedInputRow} (JSON/Kafka-style ingestion) and {@link ListBasedInputRow} + * (CSV/TSV/delimited ingestion, via {@code DelimitedValueReader}) — both expose their raw field names + * via a {@code Map}-shaped view ({@code getEvent()} / {@code asMap()}). Any other {@link InputRow} + * implementation has no generic way to enumerate its raw fields; rather than silently dropping metric + * values for such rows, this throws so the gap is caught instead of surfacing as null/0 aggregations. + */ + private static Set<String> resolveNonDimensionEventFields(final InputRow inputRow) + { + final Set<String> allFields; + if (inputRow instanceof MapBasedInputRow) { + allFields = ((MapBasedInputRow) inputRow).getEvent().keySet(); + } else if (inputRow instanceof ListBasedInputRow) { + allFields = ((ListBasedInputRow) inputRow).asMap().keySet(); + } else { + throw DruidException.defensive( + "ScanTransformer does not support input rows of type[%s]; only MapBasedInputRow and " + + "ListBasedInputRow are supported", + inputRow.getClass().getName() + ); + } + final Set<String> nonDimensionFields = new LinkedHashSet<>(allFields); + nonDimensionFields.removeAll(inputRow.getDimensions()); + return nonDimensionFields; + } + + private List<String> resolveColumnsForRow(final InputRow inputRow, final Set<String> nonDimensionEventFields) + { + final Set<String> columns = new LinkedHashSet<>(); Review Comment: P2 [P2] Honor ScanQuery.columns projection **Finding:** ScanTransformSpec accepts an embedded ScanQuery, whose columns property is the documented projection, but resolveColumnsForRow unconditionally adds every input dimension, every non-dimension event field, every query virtual column, and every unnest output. A query with columns: ["tag"] therefore still emits fields such as user, tags, and unrelated metrics, silently diverging from ScanQuery semantics and potentially storing fields the query explicitly excluded. **Suggestion:** Apply query.getColumns() when constructing the output column set, while explicitly retaining any raw fields that ingestion still requires for filters or aggregators, or reject non-empty projections until that distinction is implemented. Add coverage for a non-empty columns list. -- 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]
