CritasWang commented on code in PR #793: URL: https://github.com/apache/tsfile/pull/793#discussion_r3136400036
########## java/tools/src/main/java/org/apache/tsfile/tools/ParquetSourceReader.java: ########## @@ -0,0 +1,315 @@ +/* + * 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.tsfile.tools; + +import org.apache.tsfile.enums.TSDataType; + +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.convert.GroupRecordConverter; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.io.ColumnIOFactory; +import org.apache.parquet.io.LocalInputFile; +import org.apache.parquet.io.MessageColumnIO; +import org.apache.parquet.io.RecordReader; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class ParquetSourceReader implements SourceReader { + + private static final Logger LOGGER = LoggerFactory.getLogger(ParquetSourceReader.class); + + private final File sourceFile; + private ImportSchema schema; + private ParquetFileReader parquetReader; + private MessageType parquetSchema; + private boolean exhausted; + + private String overrideTableName; + private String overrideTimePrecision; + + public ParquetSourceReader(File sourceFile, ImportSchema schema) { + this.sourceFile = sourceFile; + this.schema = schema; + this.exhausted = false; + } + + public ParquetSourceReader(File sourceFile) { + this.sourceFile = sourceFile; + this.schema = null; + this.exhausted = false; + } + + public void setOverrideTableName(String tableName) { + this.overrideTableName = tableName; + } + + public void setOverrideTimePrecision(String timePrecision) { + this.overrideTimePrecision = timePrecision; + } + + @Override + public ImportSchema inferSchema() { + if (schema != null) { + throw new UnsupportedOperationException("inferSchema() is only available in auto mode"); + } + + try { + ensureReaderOpen(); + + List<String> columnNames = new ArrayList<>(); + List<TSDataType> columnTypes = new ArrayList<>(); + String detectedTimePrecision = null; + + for (Type field : parquetSchema.getFields()) { + String name = field.getName(); + columnNames.add(name); + + if (field.isPrimitive()) { + PrimitiveType pt = field.asPrimitiveType(); + TSDataType tsType = mapParquetType(pt); + columnTypes.add(tsType); + + if (("time".equals(name) || "TIME".equals(name)) && detectedTimePrecision == null) { + detectedTimePrecision = detectTimestampPrecision(pt); + } + } else { + columnTypes.add(TSDataType.STRING); + } + } + + String timeColumn = AutoSchemaInferer.detectTimeColumn(columnNames); + TSDataType[] types = columnTypes.toArray(new TSDataType[0]); + + String tableName = + overrideTableName != null + ? overrideTableName + : AutoSchemaInferer.deriveTableName(sourceFile.getName(), "parquet_data"); + + String timePrecision; + if (overrideTimePrecision != null) { + timePrecision = overrideTimePrecision; + } else if (detectedTimePrecision != null) { + timePrecision = detectedTimePrecision; + } else { + timePrecision = "ms"; + } + + schema = + AutoSchemaInferer.buildAutoSchema( + tableName, timeColumn, columnNames, types, timePrecision); + return schema; + } catch (IOException e) { + throw new RuntimeException("Failed to infer schema from: " + sourceFile.getAbsolutePath(), e); + } + } + + @Override + public SourceBatch readBatch() { + if (exhausted) { + return null; + } + + try { + ensureReaderOpen(); + + PageReadStore rowGroup = parquetReader.readNextRowGroup(); + if (rowGroup == null) { + exhausted = true; + return null; + } + + long rowCount = rowGroup.getRowCount(); + MessageColumnIO columnIO = new ColumnIOFactory().getColumnIO(parquetSchema); + RecordReader<Group> recordReader = + columnIO.getRecordReader(rowGroup, new GroupRecordConverter(parquetSchema)); + + List<String> schemaColumnNames = getSchemaColumnNames(); + Map<String, Integer> parquetColIndex = buildParquetColumnIndex(); + + int numCols = schemaColumnNames.size(); + List<Object[]> rows = new ArrayList<>((int) rowCount); + + for (long r = 0; r < rowCount; r++) { + Group group = recordReader.read(); + Object[] row = new Object[numCols]; + + for (int c = 0; c < numCols; c++) { + String colName = schemaColumnNames.get(c); + Integer pIdx = parquetColIndex.get(colName); + if (pIdx == null) { + row[c] = null; + continue; + } + + try { + if (group.getFieldRepetitionCount(pIdx) == 0) { + row[c] = null; + } else { + row[c] = extractValue(group, pIdx); + } + } catch (RuntimeException e) { + row[c] = null; + } + } + rows.add(row); + } + + return SourceBatch.fromRows(schemaColumnNames, rows); + } catch (IOException e) { + LOGGER.error("Error reading Parquet file: " + sourceFile.getAbsolutePath(), e); + exhausted = true; + return null; + } + } + + @Override + public void close() { + if (parquetReader != null) { + try { + parquetReader.close(); + } catch (IOException e) { + LOGGER.error("Error closing Parquet reader", e); + } + parquetReader = null; + } + } + + private void ensureReaderOpen() throws IOException { + if (parquetReader == null) { + parquetReader = ParquetFileReader.open(new LocalInputFile(sourceFile.toPath())); + parquetSchema = parquetReader.getFooter().getFileMetaData().getSchema(); + } + } + + private List<String> getSchemaColumnNames() { + List<String> names = new ArrayList<>(); + for (ImportSchema.SourceColumn col : schema.getSourceColumns()) { + if (!col.isSkip()) { Review Comment: In schema mode these readers drop skipped source columns from the SourceBatch, but TabletBuilder still indexes values by the original source_columns positions. With a schema like time, unused SKIP, value, getSchemaColumnNames() returns only [time, value], while TabletBuilder maps value to index 2, so batch.getValue(row, 2) throws and the import fails (or later columns shift incorrectly). This affects real Parquet/Arrow schemas using SKIP, and the current Parquet test only checks readBatch(), not the ImportExecutor path. Keep placeholder columns for skipped entries or change the builder mapping to use batch column names rather than schema positions. -- 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]
