jt2594838 commented on code in PR #17674: URL: https://github.com/apache/iotdb/pull/17674#discussion_r3433873740
########## iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeTsFileTabletIterator.java: ########## @@ -0,0 +1,558 @@ +/* + * 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.db.storageengine.load.converter; + +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; +import org.apache.iotdb.commons.pipe.datastructure.pattern.IoTDBTreePattern; +import org.apache.iotdb.commons.pipe.datastructure.pattern.TreePattern; +import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.query.TsFileInsertionEventQueryParser; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.scan.TsFileInsertionEventScanParser; +import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.TimeseriesMetadata; +import org.apache.tsfile.read.TsFileDeviceIterator; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.write.record.Tablet; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Load uses scan parsing first for throughput. If scan parsing hits corruption, fall back to query + * parsing for the remaining measurements and devices so later data can still be loaded. + */ +class LoadTreeTsFileTabletIterator + implements Iterable<Pair<Tablet, Boolean>>, Iterator<Pair<Tablet, Boolean>>, AutoCloseable { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadTreeTsFileTabletIterator.class); + + private static final TreePattern LOAD_TREE_PATTERN = new IoTDBTreePattern(null); + + private final File file; + private final boolean isWithMod; + private final ArrayDeque<QueryTask> pendingQueryTasks = new ArrayDeque<>(); + + private TsFileInsertionEventScanParser scanParser; + private QueryTask activeQueryTask; + private TsFileInsertionEventQueryParser activeQueryParser; + private Iterator<Pair<Tablet, Boolean>> activeIterator; + private boolean scanInitialized; + private boolean fallbackTriggered; + + private IDeviceID lastEmittedDevice; + private List<String> lastEmittedMeasurements = Collections.emptyList(); + private long lastEmittedTimestamp = Long.MIN_VALUE; + + private IDeviceID lastScanTabletDevice; + private List<String> lastScanTabletMeasurements = Collections.emptyList(); + private final Map<IDeviceID, Set<String>> fullyEmittedMeasurementsByDevice = + new LinkedHashMap<>(); + + LoadTreeTsFileTabletIterator(final File file, final boolean isWithMod) { + this.file = file; + this.isWithMod = isWithMod; + } + + @Override + public Iterator<Pair<Tablet, Boolean>> iterator() { + return this; + } + + @Override + public boolean hasNext() { + while (true) { + try { + ensureActiveIterator(); + if (Objects.isNull(activeIterator)) { + close(); + return false; + } + + if (activeIterator.hasNext()) { + return true; + } + + if (!switchToNextIterator()) { + close(); + return false; + } + } catch (final Exception e) { + if (recoverFromIteratorFailure(e)) { + continue; + } + close(); + throw toRuntimeException(e); + } + } + } + + @Override + public Pair<Tablet, Boolean> next() { + while (true) { + if (!hasNext()) { + close(); + throw new NoSuchElementException(); + } + + try { + final Pair<Tablet, Boolean> next = activeIterator.next(); + recordProgress(next); + return next; + } catch (final Exception e) { + if (recoverFromIteratorFailure(e)) { + continue; + } + close(); + throw toRuntimeException(e); + } + } + } + + private void ensureActiveIterator() throws Exception { + if (Objects.nonNull(activeIterator)) { + return; + } + + if (!scanInitialized && !fallbackTriggered) { + scanInitialized = true; + try { + scanParser = + new TsFileInsertionEventScanParser( + file, LOAD_TREE_PATTERN, Long.MIN_VALUE, Long.MAX_VALUE, null, null, isWithMod); + activeIterator = scanParser.toTabletWithIsAligneds().iterator(); + return; + } catch (final Exception e) { + if (!switchFromScanToQuery(e)) { + throw toRuntimeException(e); + } + } + } + + activateNextQueryParser(); + } + + private boolean switchToNextIterator() { + if (Objects.nonNull(activeQueryParser)) { + closeActiveQueryParser(); + return activateNextQueryParser(); + } + + closeScanParser(); + return activateNextQueryParser(); + } + + private boolean recoverFromIteratorFailure(final Exception e) { + if (shouldRethrow(e)) { + return false; + } + + if (Objects.nonNull(activeQueryTask)) { + LOGGER.warn( + "Load: Query fallback failed for device {} measurements {} in TsFile {}. " + + "Split or skip this query task and continue.", + activeQueryTask.device, + activeQueryTask.measurements, + file.getAbsolutePath(), + e); + splitOrSkipActiveQueryTask(); + return true; + } + + return switchFromScanToQuery(e); + } + + private boolean switchFromScanToQuery(final Exception e) { + if (fallbackTriggered) { + return false; + } + + fallbackTriggered = true; + final IDeviceID currentDevice = + Objects.nonNull(scanParser) ? scanParser.getCurrentDevice() : null; + final List<String> currentMeasurements = + Objects.nonNull(scanParser) ? scanParser.getCurrentMeasurements() : Collections.emptyList(); + + markLastScanMeasurementsAsCompletedIfNeeded(currentDevice, currentMeasurements); + + closeScanParser(); + + try { + pendingQueryTasks.addAll(buildQueryTasks(currentDevice, currentMeasurements)); + } catch (final Exception queryInitException) { + LOGGER.warn( + "Load: Failed to initialize query fallback for TsFile {} after scan parser failure.", + file.getAbsolutePath(), + queryInitException); + return false; + } + + LOGGER.warn( + "Load: Scan parser detected a corrupted section in TsFile {} at device {}. " + + "Switch to query parsing for remaining devices.", + file.getAbsolutePath(), + currentDevice, + e); + return true; + } + + private ArrayDeque<QueryTask> buildQueryTasks( + final IDeviceID currentDevice, final List<String> currentMeasurements) throws IOException { + final LinkedHashMap<IDeviceID, List<String>> deviceMeasurementsMap = + readDeviceMeasurementsInOrder(); + if (deviceMeasurementsMap.isEmpty()) { + return new ArrayDeque<>(); + } + + final ArrayDeque<QueryTask> tasks = new ArrayDeque<>(); + boolean includeCurrentAndFollowingDevices = + Objects.isNull(currentDevice) || !deviceMeasurementsMap.containsKey(currentDevice); + + for (final Map.Entry<IDeviceID, List<String>> entry : deviceMeasurementsMap.entrySet()) { + final IDeviceID device = entry.getKey(); + if (!includeCurrentAndFollowingDevices && device.equals(currentDevice)) { + includeCurrentAndFollowingDevices = true; + } + if (!includeCurrentAndFollowingDevices) { + continue; + } + + if (device.equals(currentDevice)) { + addCurrentDeviceQueryTasks(tasks, device, entry.getValue(), currentMeasurements); + } else { + addQueryTaskIfNecessary(tasks, device, entry.getValue(), Long.MIN_VALUE, Long.MAX_VALUE); + } + } + + return tasks; + } + + private LinkedHashMap<IDeviceID, List<String>> readDeviceMeasurementsInOrder() + throws IOException { + final LinkedHashMap<IDeviceID, List<String>> deviceMeasurementsMap = new LinkedHashMap<>(); + try (final TsFileSequenceReader reader = new TsFileSequenceReader(file.getAbsolutePath())) { + final TsFileDeviceIterator deviceIterator = reader.getAllDevicesIteratorWithIsAligned(); + while (deviceIterator.hasNext()) { + final IDeviceID device = deviceIterator.next().getLeft(); + deviceMeasurementsMap.put( + device, + reader.readDeviceMetadata(device).values().stream() + .map(TimeseriesMetadata::getMeasurementId) + .collect(Collectors.toList())); + } + } + return deviceMeasurementsMap; + } Review Comment: May use org.apache.tsfile.read.TsFileSequenceReader#getAllTimeseriesMetadata ########## iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeTsFileTabletIterator.java: ########## @@ -0,0 +1,558 @@ +/* + * 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.db.storageengine.load.converter; + +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; +import org.apache.iotdb.commons.pipe.datastructure.pattern.IoTDBTreePattern; +import org.apache.iotdb.commons.pipe.datastructure.pattern.TreePattern; +import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.query.TsFileInsertionEventQueryParser; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.scan.TsFileInsertionEventScanParser; +import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.TimeseriesMetadata; +import org.apache.tsfile.read.TsFileDeviceIterator; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.write.record.Tablet; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Load uses scan parsing first for throughput. If scan parsing hits corruption, fall back to query + * parsing for the remaining measurements and devices so later data can still be loaded. + */ +class LoadTreeTsFileTabletIterator + implements Iterable<Pair<Tablet, Boolean>>, Iterator<Pair<Tablet, Boolean>>, AutoCloseable { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadTreeTsFileTabletIterator.class); + + private static final TreePattern LOAD_TREE_PATTERN = new IoTDBTreePattern(null); + + private final File file; + private final boolean isWithMod; + private final ArrayDeque<QueryTask> pendingQueryTasks = new ArrayDeque<>(); + + private TsFileInsertionEventScanParser scanParser; + private QueryTask activeQueryTask; + private TsFileInsertionEventQueryParser activeQueryParser; + private Iterator<Pair<Tablet, Boolean>> activeIterator; + private boolean scanInitialized; + private boolean fallbackTriggered; + + private IDeviceID lastEmittedDevice; + private List<String> lastEmittedMeasurements = Collections.emptyList(); + private long lastEmittedTimestamp = Long.MIN_VALUE; + + private IDeviceID lastScanTabletDevice; + private List<String> lastScanTabletMeasurements = Collections.emptyList(); + private final Map<IDeviceID, Set<String>> fullyEmittedMeasurementsByDevice = + new LinkedHashMap<>(); + + LoadTreeTsFileTabletIterator(final File file, final boolean isWithMod) { + this.file = file; + this.isWithMod = isWithMod; + } + + @Override + public Iterator<Pair<Tablet, Boolean>> iterator() { + return this; + } + + @Override + public boolean hasNext() { + while (true) { + try { + ensureActiveIterator(); + if (Objects.isNull(activeIterator)) { + close(); + return false; + } + + if (activeIterator.hasNext()) { + return true; + } + + if (!switchToNextIterator()) { + close(); + return false; + } + } catch (final Exception e) { + if (recoverFromIteratorFailure(e)) { + continue; + } + close(); + throw toRuntimeException(e); + } + } + } + + @Override + public Pair<Tablet, Boolean> next() { + while (true) { + if (!hasNext()) { + close(); + throw new NoSuchElementException(); + } + + try { + final Pair<Tablet, Boolean> next = activeIterator.next(); + recordProgress(next); + return next; + } catch (final Exception e) { + if (recoverFromIteratorFailure(e)) { + continue; + } + close(); + throw toRuntimeException(e); + } + } + } + + private void ensureActiveIterator() throws Exception { + if (Objects.nonNull(activeIterator)) { + return; + } + + if (!scanInitialized && !fallbackTriggered) { + scanInitialized = true; + try { + scanParser = + new TsFileInsertionEventScanParser( + file, LOAD_TREE_PATTERN, Long.MIN_VALUE, Long.MAX_VALUE, null, null, isWithMod); + activeIterator = scanParser.toTabletWithIsAligneds().iterator(); + return; + } catch (final Exception e) { + if (!switchFromScanToQuery(e)) { + throw toRuntimeException(e); + } + } + } + + activateNextQueryParser(); + } + + private boolean switchToNextIterator() { + if (Objects.nonNull(activeQueryParser)) { + closeActiveQueryParser(); + return activateNextQueryParser(); + } + + closeScanParser(); + return activateNextQueryParser(); + } + + private boolean recoverFromIteratorFailure(final Exception e) { + if (shouldRethrow(e)) { + return false; + } + + if (Objects.nonNull(activeQueryTask)) { + LOGGER.warn( + "Load: Query fallback failed for device {} measurements {} in TsFile {}. " + + "Split or skip this query task and continue.", + activeQueryTask.device, + activeQueryTask.measurements, + file.getAbsolutePath(), + e); + splitOrSkipActiveQueryTask(); + return true; + } + + return switchFromScanToQuery(e); + } + + private boolean switchFromScanToQuery(final Exception e) { + if (fallbackTriggered) { + return false; + } + + fallbackTriggered = true; + final IDeviceID currentDevice = + Objects.nonNull(scanParser) ? scanParser.getCurrentDevice() : null; + final List<String> currentMeasurements = + Objects.nonNull(scanParser) ? scanParser.getCurrentMeasurements() : Collections.emptyList(); + + markLastScanMeasurementsAsCompletedIfNeeded(currentDevice, currentMeasurements); + + closeScanParser(); + + try { + pendingQueryTasks.addAll(buildQueryTasks(currentDevice, currentMeasurements)); + } catch (final Exception queryInitException) { + LOGGER.warn( + "Load: Failed to initialize query fallback for TsFile {} after scan parser failure.", + file.getAbsolutePath(), + queryInitException); + return false; + } + + LOGGER.warn( + "Load: Scan parser detected a corrupted section in TsFile {} at device {}. " + + "Switch to query parsing for remaining devices.", + file.getAbsolutePath(), + currentDevice, + e); + return true; + } + + private ArrayDeque<QueryTask> buildQueryTasks( + final IDeviceID currentDevice, final List<String> currentMeasurements) throws IOException { + final LinkedHashMap<IDeviceID, List<String>> deviceMeasurementsMap = + readDeviceMeasurementsInOrder(); + if (deviceMeasurementsMap.isEmpty()) { + return new ArrayDeque<>(); + } + + final ArrayDeque<QueryTask> tasks = new ArrayDeque<>(); + boolean includeCurrentAndFollowingDevices = + Objects.isNull(currentDevice) || !deviceMeasurementsMap.containsKey(currentDevice); + + for (final Map.Entry<IDeviceID, List<String>> entry : deviceMeasurementsMap.entrySet()) { + final IDeviceID device = entry.getKey(); + if (!includeCurrentAndFollowingDevices && device.equals(currentDevice)) { + includeCurrentAndFollowingDevices = true; + } + if (!includeCurrentAndFollowingDevices) { + continue; + } + + if (device.equals(currentDevice)) { + addCurrentDeviceQueryTasks(tasks, device, entry.getValue(), currentMeasurements); + } else { + addQueryTaskIfNecessary(tasks, device, entry.getValue(), Long.MIN_VALUE, Long.MAX_VALUE); + } + } + + return tasks; + } + + private LinkedHashMap<IDeviceID, List<String>> readDeviceMeasurementsInOrder() + throws IOException { + final LinkedHashMap<IDeviceID, List<String>> deviceMeasurementsMap = new LinkedHashMap<>(); + try (final TsFileSequenceReader reader = new TsFileSequenceReader(file.getAbsolutePath())) { + final TsFileDeviceIterator deviceIterator = reader.getAllDevicesIteratorWithIsAligned(); + while (deviceIterator.hasNext()) { + final IDeviceID device = deviceIterator.next().getLeft(); + deviceMeasurementsMap.put( + device, + reader.readDeviceMetadata(device).values().stream() + .map(TimeseriesMetadata::getMeasurementId) + .collect(Collectors.toList())); + } + } + return deviceMeasurementsMap; + } + + private void addCurrentDeviceQueryTasks( + final ArrayDeque<QueryTask> tasks, + final IDeviceID device, + final List<String> allMeasurements, + final List<String> currentMeasurements) { + final Set<String> completedMeasurements = + fullyEmittedMeasurementsByDevice.getOrDefault(device, Collections.emptySet()); + final Set<String> currentMeasurementSet = new LinkedHashSet<>(currentMeasurements); + + final List<String> currentMeasurementsToResume = new ArrayList<>(); + final List<String> remainingMeasurements = new ArrayList<>(); + for (final String measurement : allMeasurements) { + if (completedMeasurements.contains(measurement)) { + continue; + } + if (currentMeasurementSet.contains(measurement)) { + currentMeasurementsToResume.add(measurement); + } else { + remainingMeasurements.add(measurement); + } + } + + addQueryTaskIfNecessary( + tasks, + device, + currentMeasurementsToResume, + determineTaskResumeStartTime(device, currentMeasurementsToResume, Long.MIN_VALUE), + Long.MAX_VALUE); + addQueryTaskIfNecessary(tasks, device, remainingMeasurements, Long.MIN_VALUE, Long.MAX_VALUE); + } + + private boolean activateNextQueryParser() { + closeActiveQueryParser(); + + while (!pendingQueryTasks.isEmpty()) { + activeQueryTask = pendingQueryTasks.removeFirst(); + try { + activeQueryParser = + new TsFileInsertionEventQueryParser( + file, + LOAD_TREE_PATTERN, + activeQueryTask.startTime, + activeQueryTask.endTime, + activeQueryTask.toDeviceMeasurementsMap(), + isWithMod); + final Iterator<TabletInsertionEvent> tabletIterator = + activeQueryParser.toTabletInsertionEvents().iterator(); + activeIterator = + new Iterator<Pair<Tablet, Boolean>>() { + @Override + public boolean hasNext() { + return tabletIterator.hasNext(); + } + + @Override + public Pair<Tablet, Boolean> next() { + final TabletInsertionEvent event = tabletIterator.next(); + if (!(event instanceof PipeRawTabletInsertionEvent)) { + throw new IllegalStateException( + "Expected PipeRawTabletInsertionEvent but got " + event.getClass().getName()); + } + + final PipeRawTabletInsertionEvent rawTabletInsertionEvent = + (PipeRawTabletInsertionEvent) event; + return new Pair<>( + rawTabletInsertionEvent.convertToTablet(), rawTabletInsertionEvent.isAligned()); + } + }; + return true; + } catch (final Exception e) { + LOGGER.warn( + "Load: Failed to initialize query fallback for device {} measurements {} in TsFile {}. " + + "Split or skip this query task and continue.", + activeQueryTask.device, + activeQueryTask.measurements, + file.getAbsolutePath(), + e); + splitOrSkipActiveQueryTask(); + } + } + + activeIterator = null; + return false; + } + + private void recordProgress(final Pair<Tablet, Boolean> tabletWithIsAligned) { + final Tablet tablet = tabletWithIsAligned.getLeft(); + if (Objects.isNull(tablet) || tablet.getRowSize() == 0) { + return; + } + + final IDeviceID device = IDeviceID.Factory.DEFAULT_FACTORY.create(tablet.getDeviceId()); + final List<String> measurements = extractMeasurementNames(tablet); + + if (Objects.isNull(activeQueryParser)) { + recordScanProgress(device, measurements); + } + + lastEmittedDevice = device; + lastEmittedMeasurements = measurements; Review Comment: Why not directly use the measurements array from the tablet? -- 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]
