FrankChen021 commented on code in PR #19379: URL: https://github.com/apache/druid/pull/19379#discussion_r3979511840
########## processing/src/main/java/org/apache/druid/segment/transform/ScanTransformer.java: ########## @@ -0,0 +1,317 @@ +/* + * 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.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)); + } + + final List<InputRow> result = new ArrayList<>(); + while (!cursor.isDone()) { + 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)); + 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, so it falls back to an empty set. + */ + 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 { + return Set.of(); Review Comment: P2 Preserve metric fields for non-map input rows The new batch integration applies `ScanTransformSpec` to every `InputSourceReader`, but this fallback drops every raw field not listed in `getDimensions()` for any non-`MapBasedInputRow`/`ListBasedInputRow`. `DeltaInputSource` emits `DeltaInputRow`; with a metric source such as `bytes`, `ScanTransformer` therefore omits `bytes` from each expanded event and the aggregator reads null/0. Enumerate the required non-dimension fields (or reject unsupported row types) so batch inputs retain metric values. ########## processing/src/main/java/org/apache/druid/segment/transform/ScanTransformer.java: ########## @@ -0,0 +1,317 @@ +/* + * 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.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)); + } + + final List<InputRow> result = new ArrayList<>(); + while (!cursor.isDone()) { Review Comment: P2 Apply ScanQuery pagination during expansion `ScanTransformSpec` accepts a full `ScanQuery`, but this loop (and `SettableRowCursorFactory`) never uses `getScanRowsOffset()` or `getScanRowsLimit()`. Since the normal scan query toolchest applies those outside `ScanQueryEngine`, an embedded query such as `limit: 1` on a row with three unnested elements still emits all three rows into ingestion. Apply the query's offset/limit here, or reject non-default pagination so valid scan-query fields cannot silently change meaning. -- 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]
