TASK 13-14: pyflink users need geography schemas

docs/static/generated/rest_v1_sql_gateway.yml

@@ -403,6 +403,7 @@ components:
       - DESCRIPTOR
       - VARIANT
       - BITMAP
+      - GEOGRAPHY
     OpenSessionRequestBody:
       type: object
       properties:

docs/static/generated/rest_v2_sql_gateway.yml

@@ -477,6 +477,7 @@ components:
       - DESCRIPTOR
       - VARIANT
       - BITMAP
+      - GEOGRAPHY
     OpenSessionRequestBody:
       type: object
       properties:

docs/static/generated/rest_v3_sql_gateway.yml

@@ -506,6 +506,7 @@ components:
       - DESCRIPTOR
       - VARIANT
       - BITMAP
+      - GEOGRAPHY
     OpenSessionRequestBody:
       type: object
       properties:

docs/static/generated/rest_v4_sql_gateway.yml

@@ -516,6 +516,7 @@ components:
       - DESCRIPTOR
       - VARIANT
       - BITMAP
+      - GEOGRAPHY
     OpenSessionRequestBody:
       type: object
       properties:

flink-formats/flink-parquet/src/main/java/org/apache/flink/formats/parquet/row/ParquetRowDataWriter.java

@@ -119,6 +119,8 @@ private FieldWriter createWriter(LogicalType t, Type type) {
                 case BINARY:
                 case VARBINARY:
                     return new BinaryWriter();
+                case GEOGRAPHY:
+                    return new GeographyWriter();
                 case DECIMAL:
                     DecimalType decimalType = (DecimalType) t;
                     return createDecimalWriter(decimalType.getPrecision(), decimalType.getScale());
@@ -311,6 +313,23 @@ private void writeBinary(byte[] value) {
         }
     }
 
+    private class GeographyWriter implements FieldWriter {
+
+        @Override
+        public void write(RowData row, int ordinal) {
+            writeGeography(row.getGeography(ordinal).toBytes());
+        }
+
+        @Override
+        public void write(ArrayData arrayData, int ordinal) {
+            writeGeography(arrayData.getGeography(ordinal).toBytes());
+        }
+
+        private void writeGeography(byte[] value) {
+            recordConsumer.addBinary(Binary.fromReusedByteArray(value));
+        }
+    }
+
     private class IntWriter implements FieldWriter {
 
         @Override

flink-formats/flink-parquet/src/main/java/org/apache/flink/formats/parquet/utils/ParquetSchemaConverter.java

@@ -81,6 +81,10 @@ private static Type convertToParquetType(
             case VARBINARY:
                 return Types.primitive(PrimitiveType.PrimitiveTypeName.BINARY, repetition)
                         .named(name);
+            case GEOGRAPHY:
+                return Types.primitive(PrimitiveType.PrimitiveTypeName.BINARY, repetition)
+                        .as(LogicalTypeAnnotation.geographyType())
+                        .named(name);
             case DECIMAL:
                 int precision = ((DecimalType) type).getPrecision();
                 int scale = ((DecimalType) type).getScale();

flink-formats/flink-parquet/src/main/java/org/apache/flink/formats/parquet/vector/ParquetColumnarRowSplitReader.java

@@ -18,6 +18,8 @@
 
 package org.apache.flink.formats.parquet.vector;
 
+import org.apache.flink.core.fs.FileSystem;
+import org.apache.flink.formats.parquet.ParquetInputFile;
 import org.apache.flink.formats.parquet.vector.reader.ColumnReader;
 import org.apache.flink.formats.parquet.vector.type.ParquetField;
 import org.apache.flink.table.data.columnar.ColumnarRowData;
@@ -31,12 +33,12 @@
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
+import org.apache.parquet.ParquetReadOptions;
 import org.apache.parquet.column.ColumnDescriptor;
 import org.apache.parquet.column.page.PageReadStore;
 import org.apache.parquet.filter2.compat.FilterCompat;
 import org.apache.parquet.hadoop.ParquetFileReader;
 import org.apache.parquet.hadoop.metadata.BlockMetaData;
-import org.apache.parquet.hadoop.metadata.ParquetMetadata;
 import org.apache.parquet.io.ColumnIOFactory;
 import org.apache.parquet.io.MessageColumnIO;
 import org.apache.parquet.schema.GroupType;
@@ -55,9 +57,6 @@
 import static org.apache.flink.formats.parquet.vector.ParquetSplitReaderUtil.buildFieldsList;
 import static org.apache.flink.formats.parquet.vector.ParquetSplitReaderUtil.createColumnReader;
 import static org.apache.flink.formats.parquet.vector.ParquetSplitReaderUtil.createWritableColumnVector;
-import static org.apache.parquet.filter2.compat.RowGroupFilter.filterRowGroups;
-import static org.apache.parquet.format.converter.ParquetMetadataConverter.range;
-import static org.apache.parquet.hadoop.ParquetFileReader.readFooter;
 import static org.apache.parquet.hadoop.ParquetInputFormat.getFilter;
 
 /** This reader is used to read a {@link VectorizedColumnBatch} from input split. */
@@ -122,24 +121,25 @@ public ParquetColumnarRowSplitReader(
         this.utcTimestamp = utcTimestamp;
         this.selectedTypes = selectedTypes;
         this.batchSize = batchSize;
-        // then we need to apply the predicate push down filter
-        ParquetMetadata footer =
-                readFooter(conf, path, range(splitStart, splitStart + splitLength));
-        MessageType fileSchema = footer.getFileMetaData().getSchema();
         FilterCompat.Filter filter = getFilter(conf);
-        List<BlockMetaData> blocks = filterRowGroups(filter, footer.getBlocks(), fileSchema);
-
-        this.fileSchema = footer.getFileMetaData().getSchema();
+        ParquetReadOptions parquetReadOptions =
+                ParquetReadOptions.builder()
+                        .withRange(splitStart, splitStart + splitLength)
+                        .withRecordFilter(filter)
+                        .build();
+        org.apache.flink.core.fs.Path flinkPath =
+                new org.apache.flink.core.fs.Path(path.toUri().toString());
+        FileSystem fs = flinkPath.getFileSystem();
+        ParquetInputFile inputFile =
+                new ParquetInputFile(fs.open(flinkPath), fs.getFileStatus(flinkPath).getLen());
+
+        this.reader = ParquetFileReader.open(inputFile, parquetReadOptions);
+        MessageType fileSchema = reader.getFooter().getFileMetaData().getSchema();
+        this.fileSchema = fileSchema;
         this.requestedSchema = clipParquetSchema(fileSchema, selectedFieldNames, caseSensitive);
-        this.reader =
-                new ParquetFileReader(
-                        conf, footer.getFileMetaData(), path, blocks, requestedSchema.getColumns());
+        this.reader.setRequestedSchema(requestedSchema);
 
-        long totalRowCount = 0;
-        for (BlockMetaData block : blocks) {
-            totalRowCount += block.getRowCount();
-        }
-        this.totalRowCount = totalRowCount;
+        this.totalRowCount = reader.getRecordCount();
         this.nextRow = 0;
         this.rowsInBatch = 0;
         this.rowsReturned = 0;

flink-formats/flink-parquet/src/main/java/org/apache/flink/formats/parquet/vector/ParquetSplitReaderUtil.java

@@ -36,6 +36,7 @@
 import org.apache.flink.formats.parquet.vector.type.ParquetGroupField;
 import org.apache.flink.formats.parquet.vector.type.ParquetPrimitiveField;
 import org.apache.flink.table.data.DecimalData;
+import org.apache.flink.table.data.GeographyData;
 import org.apache.flink.table.data.TimestampData;
 import org.apache.flink.table.data.columnar.vector.ColumnVector;
 import org.apache.flink.table.data.columnar.vector.VectorizedColumnBatch;
@@ -169,14 +170,17 @@ public static ColumnVector createVectorFromConstant(
             case VARCHAR:
             case BINARY:
             case VARBINARY:
+            case GEOGRAPHY:
                 HeapBytesVector bsv = new HeapBytesVector(batchSize);
                 if (value == null) {
                     bsv.fillWithNulls();
                 } else {
                     bsv.fill(
-                            value instanceof byte[]
-                                    ? (byte[]) value
-                                    : value.toString().getBytes(StandardCharsets.UTF_8));
+                            value instanceof GeographyData
+                                    ? ((GeographyData) value).toBytes()
+                                    : value instanceof byte[]
+                                            ? (byte[]) value
+                                            : value.toString().getBytes(StandardCharsets.UTF_8));
                 }
                 return bsv;
             case BOOLEAN:
@@ -344,6 +348,7 @@ public static ColumnReader createColumnReader(
             case VARCHAR:
             case BINARY:
             case VARBINARY:
+            case GEOGRAPHY:
                 return new BytesColumnReader(
                         descriptors.get(0), pages.getPageReader(descriptors.get(0)));
             case TIMESTAMP_WITHOUT_TIME_ZONE:
@@ -438,6 +443,7 @@ public static WritableColumnVector createWritableColumnVector(
             case VARCHAR:
             case BINARY:
             case VARBINARY:
+            case GEOGRAPHY:
                 checkArgument(
                         typeName == PrimitiveType.PrimitiveTypeName.BINARY,
                         "Unexpected type: %s",

flink-formats/flink-parquet/src/main/java/org/apache/flink/formats/parquet/vector/reader/NestedPrimitiveColumnReader.java

@@ -238,6 +238,7 @@ private Object readPrimitiveTypedRow(LogicalType category) {
             case VARCHAR:
             case BINARY:
             case VARBINARY:
+            case GEOGRAPHY:
                 return dataColumn.readBytes();
             case BOOLEAN:
                 return dataColumn.readBoolean();
@@ -283,6 +284,7 @@ private Object dictionaryDecodeValue(LogicalType category, Integer dictionaryVal
             case VARCHAR:
             case BINARY:
             case VARBINARY:
+            case GEOGRAPHY:
                 return dictionary.readBytes(dictionaryValue);
             case DATE:
             case TIME_WITHOUT_TIME_ZONE:
@@ -324,6 +326,7 @@ private WritableColumnVector fillColumnVector(int total, List valueList) {
             case VARCHAR:
             case BINARY:
             case VARBINARY:
+            case GEOGRAPHY:
                 HeapBytesVector heapBytesVector = new HeapBytesVector(total);
                 for (int i = 0; i < valueList.size(); i++) {
                     byte[] src = ((List<byte[]>) valueList).get(i);

flink-formats/flink-parquet/src/test/java/org/apache/flink/formats/parquet/row/ParquetRowDataWriterTest.java

@@ -24,6 +24,9 @@
 import org.apache.flink.formats.parquet.ParquetWriterFactory;
 import org.apache.flink.formats.parquet.vector.ParquetColumnarRowSplitReader;
 import org.apache.flink.formats.parquet.vector.ParquetSplitReaderUtil;
+import org.apache.flink.table.data.GenericArrayData;
+import org.apache.flink.table.data.GenericRowData;
+import org.apache.flink.table.data.GeographyData;
 import org.apache.flink.table.data.RowData;
 import org.apache.flink.table.data.util.DataFormatConverters;
 import org.apache.flink.table.types.DataType;
@@ -33,6 +36,7 @@
 import org.apache.flink.table.types.logical.DecimalType;
 import org.apache.flink.table.types.logical.DoubleType;
 import org.apache.flink.table.types.logical.FloatType;
+import org.apache.flink.table.types.logical.GeographyType;
 import org.apache.flink.table.types.logical.IntType;
 import org.apache.flink.table.types.logical.MapType;
 import org.apache.flink.table.types.logical.RowType;
@@ -49,10 +53,14 @@
 import org.apache.avro.util.Utf8;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.parquet.avro.AvroParquetReader;
+import org.apache.parquet.hadoop.ParquetFileReader;
 import org.apache.parquet.hadoop.ParquetOutputFormat;
 import org.apache.parquet.hadoop.ParquetReader;
 import org.apache.parquet.hadoop.util.HadoopInputFile;
 import org.apache.parquet.io.InputFile;
+import org.apache.parquet.io.LocalInputFile;
+import org.apache.parquet.schema.LogicalTypeAnnotation;
+import org.apache.parquet.schema.PrimitiveType;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
@@ -77,6 +85,16 @@
 /** Test for {@link ParquetRowDataBuilder} and {@link ParquetRowDataWriter}. */
 class ParquetRowDataWriterTest {
 
+    private static final byte[] POINT_WKB =
+            new byte[] {
+                1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, (byte) 0xF0, 0x3F, 0, 0, 0, 0, 0, 0, 0, 0x40
+            };
+
+    private static final byte[] SECOND_POINT_WKB =
+            new byte[] {
+                1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x40, 0, 0, 0, 0, 0, 0, 8, 0x40
+            };
+
     private static final RowType ROW_TYPE =
             RowType.of(
                     new VarCharType(VarCharType.MAX_LENGTH),
@@ -192,6 +210,99 @@ public void testInt64Timestamp(@TempDir java.nio.file.Path folder) throws Except
         invalidTypeTest(folder, conf, false);
     }
 
+    @Test
+    void testGeographyType(@TempDir java.nio.file.Path folder) throws Exception {
+        RowType rowType = RowType.of(new GeographyType());
+        Path path = new Path(folder.toString(), UUID.randomUUID().toString());
+        Configuration conf = new Configuration();
+
+        ParquetWriterFactory<RowData> factory =
+                ParquetRowDataBuilder.createWriterFactory(rowType, conf, true);
+        BulkWriter<RowData> writer =
+                factory.create(path.getFileSystem().create(path, FileSystem.WriteMode.OVERWRITE));
+        writer.addElement(GenericRowData.of(GeographyData.fromBytes(POINT_WKB)));
+        writer.flush();
+        writer.finish();
+
+        try (ParquetFileReader fileReader =
+                ParquetFileReader.open(new LocalInputFile(new File(path.getPath()).toPath()))) {
+            PrimitiveType field =
+                    fileReader.getFileMetaData().getSchema().getType(0).asPrimitiveType();
+            assertThat(field.getPrimitiveTypeName())
+                    .isEqualTo(PrimitiveType.PrimitiveTypeName.BINARY);
+            assertThat(field.getLogicalTypeAnnotation())
+                    .isInstanceOf(LogicalTypeAnnotation.GeographyLogicalTypeAnnotation.class);
+        }
+
+        ParquetColumnarRowSplitReader reader =
+                ParquetSplitReaderUtil.genPartColumnarRowReader(
+                        true,
+                        true,
+                        conf,
+                        rowType.getFieldNames().toArray(new String[0]),
+                        rowType.getChildren().stream()
+                                .map(TypeConversions::fromLogicalToDataType)
+                                .toArray(DataType[]::new),
+                        new HashMap<>(),
+                        IntStream.range(0, rowType.getFieldCount()).toArray(),
+                        50,
+                        path,
+                        0,
+                        Long.MAX_VALUE);
+
+        assertThat(reader.reachedEnd()).isFalse();
+        assertThat(reader.nextRecord().getGeography(0).toBytes()).isEqualTo(POINT_WKB);
+        assertThat(reader.reachedEnd()).isTrue();
+    }
+
+    @Test
+    void testNestedGeographyType(@TempDir java.nio.file.Path folder) throws Exception {
+        RowType rowType = RowType.of(new BigIntType(), new ArrayType(true, new GeographyType()));
+        Path path = new Path(folder.toString(), UUID.randomUUID().toString());
+        Configuration conf = new Configuration();
+
+        ParquetWriterFactory<RowData> factory =
+                ParquetRowDataBuilder.createWriterFactory(rowType, conf, true);
+        BulkWriter<RowData> writer =
+                factory.create(path.getFileSystem().create(path, FileSystem.WriteMode.OVERWRITE));
+        writer.addElement(
+                GenericRowData.of(
+                        1L,
+                        new GenericArrayData(
+                                new Object[] {
+                                    GeographyData.fromBytes(POINT_WKB),
+                                    GeographyData.fromBytes(SECOND_POINT_WKB),
+                                    null
+                                })));
+        writer.flush();
+        writer.finish();
+
+        try (ParquetColumnarRowSplitReader reader =
+                ParquetSplitReaderUtil.genPartColumnarRowReader(
+                        true,
+                        true,
+                        conf,
+                        rowType.getFieldNames().toArray(new String[0]),
+                        rowType.getChildren().stream()
+                                .map(TypeConversions::fromLogicalToDataType)
+                                .toArray(DataType[]::new),
+                        new HashMap<>(),
+                        IntStream.range(0, rowType.getFieldCount()).toArray(),
+                        50,
+                        path,
+                        0,
+                        Long.MAX_VALUE)) {
+            assertThat(reader.reachedEnd()).isFalse();
+            RowData row = reader.nextRecord();
+            assertThat(row.getLong(0)).isEqualTo(1L);
+            assertThat(row.getArray(1).size()).isEqualTo(3);
+            assertThat(row.getArray(1).getGeography(0).toBytes()).isEqualTo(POINT_WKB);
+            assertThat(row.getArray(1).getGeography(1).toBytes()).isEqualTo(SECOND_POINT_WKB);
+            assertThat(row.getArray(1).isNullAt(2)).isTrue();
+            assertThat(reader.reachedEnd()).isTrue();
+        }
+    }
+
     private void innerTest(java.nio.file.Path folder, Configuration conf, boolean utcTimestamp)
             throws IOException {
         Path path = new Path(folder.toString(), UUID.randomUUID().toString());

flink-formats/pom.xml

@@ -30,7 +30,7 @@ under the License.
 	</parent>
 
 	<properties>
-		<flink.format.parquet.version>1.15.2</flink.format.parquet.version>
+		<flink.format.parquet.version>1.16.0</flink.format.parquet.version>
 	</properties>
 
 	<artifactId>flink-formats</artifactId>
@@ -92,3 +92,4 @@ under the License.
 	</profiles>
 
 </project>
+