keksmd commented on code in PR #975: URL: https://github.com/apache/incubator-graphar/pull/975#discussion_r4094044350
########## maven-projects/io-parquet/src/main/java/org/apache/graphar/io/parquet/ParquetPhysicalWriter.java: ########## @@ -0,0 +1,288 @@ +/* + * 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.graphar.io.parquet; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.time.Instant; +import java.time.LocalDate; +import java.util.List; +import java.util.Objects; +import org.apache.graphar.io.BatchCursor; +import org.apache.graphar.io.ColumnType; +import org.apache.graphar.io.Field; +import org.apache.graphar.io.PhysicalWriter; +import org.apache.graphar.io.RecordBatch; +import org.apache.graphar.io.Schema; +import org.apache.graphar.io.WriteMode; +import org.apache.graphar.io.WriteRequest; +import org.apache.graphar.storage.Storage; +import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; +import org.apache.parquet.hadoop.ParquetFileWriter; +import org.apache.parquet.hadoop.ParquetWriter; +import org.apache.parquet.hadoop.example.ExampleParquetWriter; +import org.apache.parquet.io.api.Binary; +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.apache.parquet.schema.Types; + +/** A storage-backed writer for flat, primitive Parquet batches. */ +public final class ParquetPhysicalWriter implements PhysicalWriter { + private static final int INDEXED_PAGE_ROW_COUNT = 1024; + + private final Storage storage; + + /** + * Maps the request mode onto a Parquet file mode. Parquet cannot append to a closed file + * without rewriting it, so {@link WriteMode#APPEND} is rejected instead of being downgraded to + * an overwrite. + */ + static ParquetFileWriter.Mode fileMode(WriteMode mode) { + switch (Objects.requireNonNull(mode, "mode")) { + case CREATE_NEW: + return ParquetFileWriter.Mode.CREATE; + case OVERWRITE: + return ParquetFileWriter.Mode.OVERWRITE; + default: + throw new UnsupportedOperationException( + "Parquet cannot append to an existing file; write mode " + mode); + } + } + + /** Creates a writer that resolves every output URI through {@code storage}. */ + public ParquetPhysicalWriter(Storage storage) { + this.storage = Objects.requireNonNull(storage, "storage"); + } + + @Override + public void write(WriteRequest request, BatchCursor batches) throws IOException { + Objects.requireNonNull(request, "request"); + Objects.requireNonNull(batches, "batches"); + MessageType parquetSchema = parquetSchema(request.schema()); + ParquetFileWriter.Mode mode = fileMode(request.mode()); + try (ParquetWriter<Group> writer = + ExampleParquetWriter.builder( + new ParquetOutputFile(storage.outputFile(request.uri()))) + .withType(parquetSchema) + .withWriteMode(mode) + .withPageRowCountLimit(INDEXED_PAGE_ROW_COUNT) + .build()) { + SimpleGroupFactory groups = new SimpleGroupFactory(parquetSchema); + while (batches.next()) { + RecordBatch batch = + Objects.requireNonNull(batches.batch(), "batch cursor returned null"); + requireSchema(request.schema(), batch.schema()); + for (int rowIndex = 0; rowIndex < batch.rowCount(); rowIndex++) { + writer.write(toGroup(groups, request.schema(), batch, rowIndex)); + } + } + } finally { + batches.close(); + } + } + + private static MessageType parquetSchema(Schema schema) { + Types.MessageTypeBuilder builder = Types.buildMessage(); + for (Field field : schema.fields()) { + if (field.type().kind() == ColumnType.Kind.LIST) { + builder.addField(listType(field)); + continue; + } + Types.PrimitiveBuilder<Types.GroupBuilder<MessageType>> primitive = + builder.primitive(physicalType(field.type()), repetition(field)); + LogicalTypeAnnotation logicalType = logicalType(field.type()); + if (logicalType != null) { + primitive.as(logicalType); + } + primitive.named(field.name()); + } + return builder.named("graphar"); + } + + private static Type listType(Field field) { + ColumnType element = field.type().elementType().orElseThrow(); + Type value = listElementType(element); + return (field.nullable() ? Types.optionalList() : Types.requiredList()) + .element(value) + .named(field.name()); + } + + private static Type listElementType(ColumnType element) { + LogicalTypeAnnotation logical = logicalType(element); + if (logical == null) { + return Types.repeated(physicalType(element)).named("element"); Review Comment: Changed together with the reader, see the reply there. ########## maven-projects/io-parquet/src/main/java/org/apache/graphar/io/parquet/ParquetBatchCursor.java: ########## @@ -0,0 +1,350 @@ +/* + * 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.graphar.io.parquet; + +import java.io.IOException; +import java.time.Instant; +import java.time.LocalDate; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.graphar.io.BatchCursor; +import org.apache.graphar.io.ReadRequest; +import org.apache.graphar.io.RecordBatch; +import org.apache.graphar.io.RowRange; +import org.apache.graphar.io.Schema; +import org.apache.graphar.io.ValueVector; +import org.apache.graphar.io.VectorRecordBatch; +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.filter2.columnindex.RowRanges; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore.MissingOffsetIndexException; +import org.apache.parquet.io.ColumnIOFactory; +import org.apache.parquet.io.MessageColumnIO; +import org.apache.parquet.io.RecordReader; +import org.apache.parquet.schema.MessageType; + +/** Streams materialized Parquet row groups as neutral record batches. */ +final class ParquetBatchCursor implements BatchCursor { + private static final int BATCH_ROWS = 1_024; + private final ParquetFileReader fileReader; + private final MessageType fileSchema; + private final MessageType readSchema; + private final List<ParquetColumn> readColumns; + private final Schema outputSchema; + private final Map<String, Integer> readColumnIndexes; + private final int[] outputIndexes; + private final long rangeStart; + private final long rangeEnd; + private final long limit; + private final List<BlockRange> rowGroups; + private int nextRowGroup; + private long emitted; + private long rowsRemainingInGroup; + private PageReadStore pages; + private RecordReader<Group> rows; + private boolean exhausted; + private boolean closed; + private RecordBatch current; + + ParquetBatchCursor( + ParquetFileReader fileReader, + MessageType fileSchema, + MessageType readSchema, + List<ParquetColumn> readColumns, + List<ParquetColumn> outputColumns, + Schema outputSchema, + ReadRequest request) { + this.fileReader = fileReader; + this.fileSchema = fileSchema; + this.readSchema = readSchema; + this.readColumns = readColumns; + this.outputSchema = outputSchema; + this.readColumnIndexes = indexes(readColumns); + this.outputIndexes = outputIndexes(outputColumns, readColumnIndexes); + RowRange range = request.rowRange().orElse(null); + this.rangeStart = range == null ? 0 : range.startInclusive(); + this.rangeEnd = range == null ? Long.MAX_VALUE : range.endExclusive(); + this.limit = request.limit().isPresent() ? request.limit().getAsLong() : Long.MAX_VALUE; + this.rowGroups = rowGroups(fileReader.getRowGroups()); + } + + @Override + public boolean next() throws IOException { + if (closed || exhausted) { + current = null; + return false; + } + if (emitted == limit) { + finish(); + return false; + } + try { + while (true) { + if (rows == null && !openNextRowGroup()) { + finish(); + return false; + } + int batchSize = + (int) Math.min(Math.min(rowsRemainingInGroup, BATCH_ROWS), limit - emitted); + Object[][] columns = new Object[outputIndexes.length][batchSize]; + for (int index = 0; index < batchSize; index++) { + Group group = rows.read(); + project(values(group), columns, index); + emitted++; + rowsRemainingInGroup--; + } + if (rowsRemainingInGroup == 0) { + closePages(); + } + current = batch(columns, batchSize); + if (emitted == limit) exhausted = true; + return true; + } + } catch (MissingOffsetIndexException exception) { + try { + closeReader(); + } catch (IOException closeException) { + exception.addSuppressed(closeException); + } + throw new UnsupportedOperationException( + "Physical Parquet row ranges require an Offset Index; refusing JVM fallback.", + exception); + } catch (IOException | RuntimeException exception) { + try { + closeReader(); + } catch (IOException closeException) { + exception.addSuppressed(closeException); + } + throw exception; + } + } + + @Override + public RecordBatch batch() { + if (current == null) { + throw new IllegalStateException("No current batch. Call next() before batch()."); + } + return current; + } + + @Override + public void close() throws IOException { + current = null; + exhausted = true; + closeReader(); + } + + private Object[] values(Group group) { + Object[] values = new Object[readColumns.size()]; + for (int index = 0; index < readColumns.size(); index++) { + if (group.getFieldRepetitionCount(index) != 0) { + values[index] = value(group, index, readColumns.get(index)); + } + } + return values; + } + + private static Object value(Group group, int index, ParquetColumn column) { + if (column.field().type().kind() == org.apache.graphar.io.ColumnType.Kind.LIST) { + return listValue(group, index, column); + } + return scalarValue(group, index, column.field().type().kind()); + } + + private static List<Object> listValue(Group group, int index, ParquetColumn column) { + Group list = group.getGroup(index, 0); + if (list.getType().getFieldCount() != 1) { + throw new IllegalArgumentException( + "Unsupported Parquet LIST field: " + column.field().name()); + } + int count = list.getFieldRepetitionCount(0); + if (count == 0) { + return List.of(); + } + List<Object> values = new ArrayList<>(count); + org.apache.graphar.io.ColumnType element = + column.field().type().elementType().orElseThrow(); + for (int elementIndex = 0; elementIndex < count; elementIndex++) { + Group elementGroup = list.getGroup(0, elementIndex); + if (elementGroup.getType().getFieldCount() != 1 + || elementGroup.getFieldRepetitionCount(0) != 1) { + throw new IllegalArgumentException( + "Unsupported Parquet LIST element: " + column.field().name()); + } + values.add(scalarValue(elementGroup, 0, element, 0)); + } + return Collections.unmodifiableList(values); + } + + private static Object scalarValue( + Group group, int index, org.apache.graphar.io.ColumnType.Kind kind) { + return scalarValue(group, index, org.apache.graphar.io.ColumnType.of(kind), 0); + } + + private static Object scalarValue( + Group group, int index, org.apache.graphar.io.ColumnType type, int repetitionIndex) { + switch (type.kind()) { + case BOOLEAN: + return group.getBoolean(index, repetitionIndex); + case INT8: + return (byte) group.getInteger(index, repetitionIndex); + case INT16: + return (short) group.getInteger(index, repetitionIndex); + case INT32: + return group.getInteger(index, repetitionIndex); + case INT64: + return group.getLong(index, repetitionIndex); + case FLOAT32: + return group.getFloat(index, repetitionIndex); + case FLOAT64: + return group.getDouble(index, repetitionIndex); + case STRING: + return group.getBinary(index, repetitionIndex).toStringUsingUTF8(); + case BINARY: + return group.getBinary(index, repetitionIndex).getBytes(); + case DATE: + return LocalDate.ofEpochDay(group.getInteger(index, repetitionIndex)); + case TIMESTAMP_MILLIS: + return Instant.ofEpochMilli(group.getLong(index, repetitionIndex)); + default: + throw new IllegalArgumentException( + "Unsupported Parquet column type: " + type.kind()); + } + } + + private void project(Object[] values, Object[][] columns, int row) { + for (int index = 0; index < outputIndexes.length; index++) { + Object value = values[outputIndexes[index]]; + columns[index][row] = value instanceof byte[] ? ((byte[]) value).clone() : value; + } + } + + private RecordBatch batch(Object[][] columns, int rowCount) { + List<ValueVector> vectors = new ArrayList<>(columns.length); + for (int index = 0; index < columns.length; index++) { + vectors.add(new ParquetValueVector(outputSchema.fields().get(index), columns[index])); + } + return new VectorRecordBatch(outputSchema, vectors, rowCount); + } + + private void finish() throws IOException { + exhausted = true; + current = null; + closeReader(); + } + + private void closeReader() throws IOException { + if (!closed) { + closed = true; + try { + closePages(); + } finally { + fileReader.close(); + } + } + } + + private boolean openNextRowGroup() throws IOException { + BlockRange rowGroup = nextRange(); + if (rowGroup == null) return false; + pages = Review Comment: Changed in eb0ad12: a range that covers the whole row group now calls `readRowGroup(index)`, and only a partial range goes through `readFilteredRowGroup`; the error message now says a partial range needs the Offset Index. For the record, parquet-java 1.18 already short-circuits `readFilteredRowGroup` to `internalReadRowGroup` when the ranges cover the block, so plain scans did not fail in practice, but relying on that was implicit. ########## maven-projects/io-parquet/src/main/java/org/apache/graphar/io/parquet/ParquetPhysicalReader.java: ########## @@ -0,0 +1,426 @@ +/* + * 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.graphar.io.parquet; + +import java.io.IOException; +import java.net.URI; +import java.util.ArrayList; +import java.util.Collections; +import java.util.EnumSet; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import org.apache.graphar.io.ColumnRef; +import org.apache.graphar.io.ColumnType; +import org.apache.graphar.io.Field; +import org.apache.graphar.io.PhysicalReader; +import org.apache.graphar.io.ReadCapability; +import org.apache.graphar.io.ReadReport; +import org.apache.graphar.io.ReadRequest; +import org.apache.graphar.io.ReadResult; +import org.apache.graphar.io.Schema; +import org.apache.graphar.storage.InputFile; +import org.apache.graphar.storage.Storage; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.io.SeekableInputStream; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; +import org.apache.parquet.schema.Type; + +/** A storage-backed reader for GraphAr primitive and LIST Parquet fields. */ +public final class ParquetPhysicalReader implements PhysicalReader { + private static final Set<ReadCapability> CAPABILITIES = + Collections.unmodifiableSet( + EnumSet.of( + ReadCapability.PROJECTION, + ReadCapability.ROW_RANGE, + ReadCapability.LIMIT)); + + private static final int DEFAULT_FOOTER_CACHE_CAPACITY = 256; + + private final Storage storage; + private final FooterCache footers; + + /** + * Creates a reader that resolves each request URI through {@code storage} and remembers the + * footers of recently read files. + */ + public ParquetPhysicalReader(Storage storage) { + this(storage, DEFAULT_FOOTER_CACHE_CAPACITY); + } + + /** + * Creates a reader that keeps at most {@code footerCacheCapacity} Parquet footers in memory. + * GraphAr chunk files are immutable once published, so repeated range reads of one chunk parse + * its footer once. A capacity of zero parses the footer on every request. + * + * @param storage resolves each request URI to a readable file + * @param footerCacheCapacity maximum number of remembered footers + */ + public ParquetPhysicalReader(Storage storage, int footerCacheCapacity) { + this.storage = Objects.requireNonNull(storage, "storage"); + this.footers = new FooterCache(footerCacheCapacity); + } + + @Override + public Set<ReadCapability> capabilities() { + return CAPABILITIES; + } + + @Override + public ReadResult read(ReadRequest request) throws IOException { + Objects.requireNonNull(request, "request"); + if (!request.filters().isEmpty()) { + throw new UnsupportedOperationException( + "Parquet filter pushdown is not implemented; refusing semantic fallback."); + } + InputFile inputFile = + Objects.requireNonNull(storage.inputFile(request.uri()), "storage inputFile"); + ParquetFileReader fileReader = null; + try { + fileReader = open(request.uri(), new ParquetInputFile(inputFile)); + MessageType fileSchema = fileReader.getFooter().getFileMetaData().getSchema(); + List<ParquetColumn> fileColumns = columns(fileSchema); + Map<String, ParquetColumn> columnsByName = byName(fileColumns); + List<ParquetColumn> outputColumns = outputColumns(request, fileColumns, columnsByName); + List<ParquetColumn> readColumns = readColumns(fileColumns, outputColumns); + MessageType readSchema = + new MessageType(fileSchema.getName(), parquetTypes(readColumns)); + fileReader.setRequestedSchema(readSchema); + + Schema outputSchema = new Schema(fields(outputColumns)); + ReadReport report = new ReadReport(applied(request), declined(request)); + ParquetBatchCursor cursor = + new ParquetBatchCursor( + fileReader, + fileSchema, + readSchema, + readColumns, + outputColumns, + outputSchema, + request); + fileReader = null; + return new ReadResult(request, cursor, report); + } finally { + if (fileReader != null) { + fileReader.close(); + } + } + } + + /** + * Opens a Parquet reader, reusing a remembered footer when the file is unchanged in size. The + * returned reader owns the stream opened here and closes it. + */ + private ParquetFileReader open(URI uri, ParquetInputFile file) throws IOException { + ParquetReadOptions options = ParquetReadOptions.builder().build(); + long size = file.getLength(); + ParquetMetadata remembered = footers.get(uri, size); Review Comment: Fixed in eb0ad12. The cache is now opt-in: `new ParquetPhysicalReader(storage)` parses the footer every time, and the capacity constructor documents that the caller guarantees no file is rewritten in place while the reader is in use. Storage exposes no version or ETag signal today, so a stronger key would have to come from a storage-api change; I kept that out of this PR. ########## maven-projects/io-parquet/src/main/java/org/apache/graphar/io/parquet/ParquetBatchCursor.java: ########## @@ -0,0 +1,350 @@ +/* + * 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.graphar.io.parquet; + +import java.io.IOException; +import java.time.Instant; +import java.time.LocalDate; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.graphar.io.BatchCursor; +import org.apache.graphar.io.ReadRequest; +import org.apache.graphar.io.RecordBatch; +import org.apache.graphar.io.RowRange; +import org.apache.graphar.io.Schema; +import org.apache.graphar.io.ValueVector; +import org.apache.graphar.io.VectorRecordBatch; +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.filter2.columnindex.RowRanges; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore.MissingOffsetIndexException; +import org.apache.parquet.io.ColumnIOFactory; +import org.apache.parquet.io.MessageColumnIO; +import org.apache.parquet.io.RecordReader; +import org.apache.parquet.schema.MessageType; + +/** Streams materialized Parquet row groups as neutral record batches. */ +final class ParquetBatchCursor implements BatchCursor { + private static final int BATCH_ROWS = 1_024; + private final ParquetFileReader fileReader; + private final MessageType fileSchema; + private final MessageType readSchema; + private final List<ParquetColumn> readColumns; + private final Schema outputSchema; + private final Map<String, Integer> readColumnIndexes; + private final int[] outputIndexes; + private final long rangeStart; + private final long rangeEnd; + private final long limit; + private final List<BlockRange> rowGroups; + private int nextRowGroup; + private long emitted; + private long rowsRemainingInGroup; + private PageReadStore pages; + private RecordReader<Group> rows; + private boolean exhausted; + private boolean closed; + private RecordBatch current; + + ParquetBatchCursor( + ParquetFileReader fileReader, + MessageType fileSchema, + MessageType readSchema, + List<ParquetColumn> readColumns, + List<ParquetColumn> outputColumns, + Schema outputSchema, + ReadRequest request) { + this.fileReader = fileReader; + this.fileSchema = fileSchema; + this.readSchema = readSchema; + this.readColumns = readColumns; + this.outputSchema = outputSchema; + this.readColumnIndexes = indexes(readColumns); + this.outputIndexes = outputIndexes(outputColumns, readColumnIndexes); + RowRange range = request.rowRange().orElse(null); + this.rangeStart = range == null ? 0 : range.startInclusive(); + this.rangeEnd = range == null ? Long.MAX_VALUE : range.endExclusive(); + this.limit = request.limit().isPresent() ? request.limit().getAsLong() : Long.MAX_VALUE; + this.rowGroups = rowGroups(fileReader.getRowGroups()); + } + + @Override + public boolean next() throws IOException { + if (closed || exhausted) { + current = null; + return false; + } + if (emitted == limit) { Review Comment: Fixed in eb0ad12. Reaching the limit now closes the file reader immediately while keeping the last batch retrievable; the top-of-method check only handles `limit == 0`. Test `releasesTheFileAsSoonAsTheLimitIsReached` asserts every opened stream is closed right after the final `next()` returning true. ########## maven-projects/io-parquet/src/main/java/org/apache/graphar/io/parquet/ParquetBatchCursor.java: ########## @@ -0,0 +1,350 @@ +/* + * 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.graphar.io.parquet; + +import java.io.IOException; +import java.time.Instant; +import java.time.LocalDate; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.graphar.io.BatchCursor; +import org.apache.graphar.io.ReadRequest; +import org.apache.graphar.io.RecordBatch; +import org.apache.graphar.io.RowRange; +import org.apache.graphar.io.Schema; +import org.apache.graphar.io.ValueVector; +import org.apache.graphar.io.VectorRecordBatch; +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.filter2.columnindex.RowRanges; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.internal.filter2.columnindex.ColumnIndexStore.MissingOffsetIndexException; +import org.apache.parquet.io.ColumnIOFactory; +import org.apache.parquet.io.MessageColumnIO; +import org.apache.parquet.io.RecordReader; +import org.apache.parquet.schema.MessageType; + +/** Streams materialized Parquet row groups as neutral record batches. */ +final class ParquetBatchCursor implements BatchCursor { + private static final int BATCH_ROWS = 1_024; + private final ParquetFileReader fileReader; + private final MessageType fileSchema; + private final MessageType readSchema; + private final List<ParquetColumn> readColumns; + private final Schema outputSchema; + private final Map<String, Integer> readColumnIndexes; + private final int[] outputIndexes; + private final long rangeStart; + private final long rangeEnd; + private final long limit; + private final List<BlockRange> rowGroups; + private int nextRowGroup; + private long emitted; + private long rowsRemainingInGroup; + private PageReadStore pages; + private RecordReader<Group> rows; + private boolean exhausted; + private boolean closed; + private RecordBatch current; + + ParquetBatchCursor( + ParquetFileReader fileReader, + MessageType fileSchema, + MessageType readSchema, + List<ParquetColumn> readColumns, + List<ParquetColumn> outputColumns, + Schema outputSchema, + ReadRequest request) { + this.fileReader = fileReader; + this.fileSchema = fileSchema; + this.readSchema = readSchema; + this.readColumns = readColumns; + this.outputSchema = outputSchema; + this.readColumnIndexes = indexes(readColumns); + this.outputIndexes = outputIndexes(outputColumns, readColumnIndexes); + RowRange range = request.rowRange().orElse(null); + this.rangeStart = range == null ? 0 : range.startInclusive(); + this.rangeEnd = range == null ? Long.MAX_VALUE : range.endExclusive(); + this.limit = request.limit().isPresent() ? request.limit().getAsLong() : Long.MAX_VALUE; + this.rowGroups = rowGroups(fileReader.getRowGroups()); + } + + @Override + public boolean next() throws IOException { + if (closed || exhausted) { + current = null; + return false; + } + if (emitted == limit) { + finish(); + return false; + } + try { + while (true) { + if (rows == null && !openNextRowGroup()) { + finish(); + return false; + } + int batchSize = + (int) Math.min(Math.min(rowsRemainingInGroup, BATCH_ROWS), limit - emitted); + Object[][] columns = new Object[outputIndexes.length][batchSize]; + for (int index = 0; index < batchSize; index++) { + Group group = rows.read(); + project(values(group), columns, index); + emitted++; + rowsRemainingInGroup--; + } + if (rowsRemainingInGroup == 0) { + closePages(); + } + current = batch(columns, batchSize); + if (emitted == limit) exhausted = true; + return true; + } + } catch (MissingOffsetIndexException exception) { + try { + closeReader(); + } catch (IOException closeException) { + exception.addSuppressed(closeException); + } + throw new UnsupportedOperationException( + "Physical Parquet row ranges require an Offset Index; refusing JVM fallback.", + exception); + } catch (IOException | RuntimeException exception) { + try { + closeReader(); + } catch (IOException closeException) { + exception.addSuppressed(closeException); + } + throw exception; + } + } + + @Override + public RecordBatch batch() { + if (current == null) { + throw new IllegalStateException("No current batch. Call next() before batch()."); + } + return current; + } + + @Override + public void close() throws IOException { + current = null; + exhausted = true; + closeReader(); + } + + private Object[] values(Group group) { + Object[] values = new Object[readColumns.size()]; + for (int index = 0; index < readColumns.size(); index++) { + if (group.getFieldRepetitionCount(index) != 0) { + values[index] = value(group, index, readColumns.get(index)); + } + } + return values; + } + + private static Object value(Group group, int index, ParquetColumn column) { + if (column.field().type().kind() == org.apache.graphar.io.ColumnType.Kind.LIST) { + return listValue(group, index, column); + } + return scalarValue(group, index, column.field().type().kind()); + } + + private static List<Object> listValue(Group group, int index, ParquetColumn column) { + Group list = group.getGroup(index, 0); + if (list.getType().getFieldCount() != 1) { + throw new IllegalArgumentException( + "Unsupported Parquet LIST field: " + column.field().name()); + } + int count = list.getFieldRepetitionCount(0); + if (count == 0) { + return List.of(); + } + List<Object> values = new ArrayList<>(count); + org.apache.graphar.io.ColumnType element = + column.field().type().elementType().orElseThrow(); + for (int elementIndex = 0; elementIndex < count; elementIndex++) { + Group elementGroup = list.getGroup(0, elementIndex); Review Comment: Renamed in eb0ad12: `listGroup` for the LIST group, `repeatedGroup` for each repeated `list` wrapper, `elementType` for the element. -- 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]
