From 600bdc798436a253edf29327964bbe7d7e1dde44 Mon Sep 17 00:00:00 2001 From: morningman Date: Mon, 21 Sep 2026 14:38:22 +0800 Subject: [PATCH 1/3] [test](arrow-flight) Pin the Doris-to-Arrow type mapping as it stands, one row per PrimitiveType The schema in the table_schema column of Flight SQL GetTables is the only statement FE makes about the Arrow type of a column, and a client that types its columns from it reads the batches BE emits as that type. Until now only the nested-type rules and DATEV2 were covered; the scalar table itself, including the cells known to disagree with BE (TIMESTAMPTZ with a literal "UTC" zone, Null for TIMEV2 / VARBINARY / AGG_STATE, see #67577), could change without any test noticing. This records the table exactly as it is, one row per PrimitiveType plus one per precision / scale band where the band picks the Arrow unit, and fails when a PrimitiveType has no row, so a new type gets a deliberate mapping rather than the default. It is the baseline for extracting the mapping into a class of its own with no value changed. Co-Authored-By: Claude Opus 5 (1M context) --- .../FlightSqlSchemaHelperArrowTypeTest.java | 122 ++++++++++++++++++ 1 file changed, 122 insertions(+) diff --git a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperArrowTypeTest.java b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperArrowTypeTest.java index d7ffadfc121015..8b74a1093ef549 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperArrowTypeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperArrowTypeTest.java @@ -17,6 +17,7 @@ package org.apache.doris.arrowflight; +import org.apache.doris.catalog.PrimitiveType; import org.apache.doris.common.jmockit.Deencapsulation; import org.apache.doris.thrift.TColumnDesc; import org.apache.doris.thrift.TPrimitiveType; @@ -26,19 +27,27 @@ import org.apache.arrow.vector.ipc.ReadChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.DateUnit; +import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import java.io.ByteArrayInputStream; import java.io.IOException; +import java.lang.reflect.Method; import java.nio.channels.Channels; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.Stream; /** * What {@code CommandGetTables} says a column is, against what the query that follows actually @@ -78,6 +87,119 @@ private static Field buildField(TColumnDesc columnDesc) { return Deencapsulation.invoke(FlightSqlSchemaHelper.class, "buildField", DB, TABLE, columnDesc); } + private static ArrowType arrowType(PrimitiveType type, Integer precision, Integer scale) throws Exception { + // Plain reflection rather than Deencapsulation: the descriptor of a column that has no precision + // or scale hands the mapping null for both, and that null is part of what the table pins. + Method getArrowType = FlightSqlSchemaHelper.class.getDeclaredMethod("getArrowType", + PrimitiveType.class, Integer.class, Integer.class); + getArrowType.setAccessible(true); + return (ArrowType) getArrowType.invoke(null, type, precision, scale); + } + + private static Arguments row(PrimitiveType type, ArrowType expected) { + return Arguments.of(type, null, null, expected); + } + + private static Arguments row(PrimitiveType type, int precision, int scale, ArrowType expected) { + return Arguments.of(type, precision, scale, expected); + } + + private static ArrowType timestamp(TimeUnit unit, String timezone) { + return new ArrowType.Timestamp(unit, timezone); + } + + /** + * The mapping as it stands today, one row per {@link PrimitiveType} (plus one per precision / scale + * band where the band picks the Arrow type). This is a record of the present, not of the ideal: the + * rows marked "kept as is" are known to disagree with what BE emits and stay that way on purpose + * until the BE Arrow type layer is reworked, after which the whole table is corrected in one step + * against a golden shared with BE (tracked in #67577). Until then a change to any row is a + * behaviour change that every {@code GetTables} client sees, and this test is what makes it + * deliberate. + */ + private static Stream mapping() { + return Stream.of( + row(PrimitiveType.BOOLEAN, new ArrowType.Bool()), + row(PrimitiveType.TINYINT, new ArrowType.Int(8, true)), + row(PrimitiveType.SMALLINT, new ArrowType.Int(16, true)), + row(PrimitiveType.INT, new ArrowType.Int(32, true)), + row(PrimitiveType.BIGINT, new ArrowType.Int(64, true)), + row(PrimitiveType.FLOAT, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + row(PrimitiveType.DOUBLE, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + // BE writes a LARGEINT as its decimal text: it does not fit decimal128. + row(PrimitiveType.LARGEINT, new ArrowType.Utf8()), + row(PrimitiveType.CHAR, new ArrowType.Utf8()), + row(PrimitiveType.VARCHAR, new ArrowType.Utf8()), + row(PrimitiveType.STRING, new ArrowType.Utf8()), + row(PrimitiveType.JSONB, new ArrowType.Utf8()), + row(PrimitiveType.VARIANT, new ArrowType.Utf8()), + // IPV4 rides in an int32 (parquet has no uint32); IPV6 is text. + row(PrimitiveType.IPV4, new ArrowType.Int(32, true)), + row(PrimitiveType.IPV6, new ArrowType.Utf8()), + // The v1 date types stay text; DATEV2 is a day number, see dateV2IsDescribedAsDate32. + row(PrimitiveType.DATE, new ArrowType.Utf8()), + row(PrimitiveType.DATETIME, new ArrowType.Utf8()), + row(PrimitiveType.DATEV2, new ArrowType.Date(DateUnit.DAY)), + // DATETIMEV2 is a wall-clock value: a timezone-naive timestamp whose unit follows the scale, + // with the bands' edges at scale 0 / 1 and 3 / 4. + row(PrimitiveType.DATETIMEV2, 18, 0, timestamp(TimeUnit.SECOND, null)), + row(PrimitiveType.DATETIMEV2, 19, 1, timestamp(TimeUnit.MILLISECOND, null)), + row(PrimitiveType.DATETIMEV2, 21, 3, timestamp(TimeUnit.MILLISECOND, null)), + row(PrimitiveType.DATETIMEV2, 22, 4, timestamp(TimeUnit.MICROSECOND, null)), + row(PrimitiveType.DATETIMEV2, 24, 6, timestamp(TimeUnit.MICROSECOND, null)), + row(PrimitiveType.TIMESTAMP_NS, 27, 9, timestamp(TimeUnit.NANOSECOND, null)), + // The same bands as DATETIMEV2, but the timezone is the literal "UTC" where BE stamps + // the session timezone. Kept as is (#67577). + row(PrimitiveType.TIMESTAMPTZ, 18, 0, timestamp(TimeUnit.SECOND, "UTC")), + row(PrimitiveType.TIMESTAMPTZ, 21, 3, timestamp(TimeUnit.MILLISECOND, "UTC")), + row(PrimitiveType.TIMESTAMPTZ, 24, 6, timestamp(TimeUnit.MICROSECOND, "UTC")), + // DECIMALV2 is always (27, 9) whatever the column declares; the v3 decimals carry theirs. + row(PrimitiveType.DECIMALV2, 10, 2, new ArrowType.Decimal(27, 9, 128)), + row(PrimitiveType.DECIMAL32, 9, 2, new ArrowType.Decimal(9, 2, 128)), + row(PrimitiveType.DECIMAL64, 18, 4, new ArrowType.Decimal(18, 4, 128)), + row(PrimitiveType.DECIMAL128, 38, 10, new ArrowType.Decimal(38, 10, 128)), + row(PrimitiveType.DECIMAL256, 76, 20, new ArrowType.Decimal(76, 20, 256)), + row(PrimitiveType.HLL, new ArrowType.Binary()), + row(PrimitiveType.BITMAP, new ArrowType.Binary()), + row(PrimitiveType.QUANTILE_STATE, new ArrowType.Binary()), + // BE emits float64 for TIMEV2 and binary for VARBINARY and AGG_STATE; the schema says + // Null for all three. Kept as is (#67577). + row(PrimitiveType.TIMEV2, 18, 0, new ArrowType.Null()), + row(PrimitiveType.VARBINARY, new ArrowType.Null()), + row(PrimitiveType.AGG_STATE, new ArrowType.Null()), + // The element types of a complex column live in the field's children, not in its type. + row(PrimitiveType.ARRAY, new ArrowType.List()), + row(PrimitiveType.MAP, new ArrowType.Map(false)), + row(PrimitiveType.STRUCT, new ArrowType.Struct()), + // Types that never name a stored column fall through to Null. + row(PrimitiveType.INVALID_TYPE, new ArrowType.Null()), + row(PrimitiveType.UNSUPPORTED, new ArrowType.Null()), + row(PrimitiveType.NULL_TYPE, new ArrowType.Null()), + row(PrimitiveType.LAMBDA_FUNCTION, new ArrowType.Null()), + row(PrimitiveType.TEMPLATE, new ArrowType.Null()), + row(PrimitiveType.BINARY, new ArrowType.Null())); + } + + @ParameterizedTest(name = "{0}({1}, {2}) is described as {3}") + @MethodSource("mapping") + public void everyTypeIsDescribedAsToday(PrimitiveType type, Integer precision, Integer scale, + ArrowType expected) throws Exception { + Assertions.assertEquals(expected, arrowType(type, precision, scale)); + } + + /** + * A type this table does not know is a type whose schema nobody has looked at: a new + * {@link PrimitiveType} must get a row here, and the row must say what BE emits for it. + */ + @Test + public void everyPrimitiveTypeHasARow() { + Set covered = mapping().map(row -> (PrimitiveType) row.get()[0]) + .collect(Collectors.toSet()); + for (PrimitiveType type : PrimitiveType.values()) { + Assertions.assertTrue(covered.contains(type), type + " has no row in the mapping table"); + } + } + /** * BE writes a DATEV2 column as {@code arrow::Date32Type} -- a day number. {@link DateUnit#MILLISECOND} * is date64, a different width and a different meaning, and a client that believes it renders and From 2371c923f6123060f3f8b79c6243d263f1ec899c Mon Sep 17 00:00:00 2001 From: morningman Date: Mon, 21 Sep 2026 14:42:19 +0800 Subject: [PATCH 2/3] [refactor](arrow-flight) Extract the Doris-to-Arrow type mapping into DorisArrowTypeMapping FE is about to hand clients more Arrow schemas than the one in GetTables: the result and parameter schemas of prepared statements, ExecuteSchema. Each of them must describe a column exactly as the others do, or a client that types its columns from one of them fails to read the batches another promised. Until now the only mapping lived as private methods of FlightSqlSchemaHelper, where a second caller could only copy it. This moves getArrowType, the column-descriptor overload and the field builder with its nested-type rules into org.apache.doris.arrow. DorisArrowTypeMapping, unchanged: the switch, the children rules and the Flight SQL column metadata are byte for byte what they were, and the table-driven test recorded in the previous commit passes against the new class with no row touched. The cells known to disagree with BE (the literal "UTC" zone on TIMESTAMPTZ, Null for TIMEV2 / VARBINARY / AGG_STATE) are carried over as they are: they are fixed in one step against a golden shared with BE once the BE Arrow type layer has been reworked, not one call site at a time (#67577). FlightSqlSchemaHelper keeps the GetTables plumbing and the schema serialization; the tests split the same way, the mapping and nesting cases into DorisArrowTypeMappingTest, the serialized-schema round trips into FlightSqlSchemaHelperSerializedSchemaTest. Co-Authored-By: Claude Opus 5 (1M context) --- .../doris/arrow/DorisArrowTypeMapping.java | 229 ++++++++++++++++++ .../arrowflight/FlightSqlSchemaHelper.java | 183 +------------- .../DorisArrowTypeMappingTest.java} | 83 ++----- ...htSqlSchemaHelperSerializedSchemaTest.java | 102 ++++++++ 4 files changed, 347 insertions(+), 250 deletions(-) create mode 100644 fe/fe-core/src/main/java/org/apache/doris/arrow/DorisArrowTypeMapping.java rename fe/fe-core/src/test/java/org/apache/doris/{arrowflight/FlightSqlSchemaHelperArrowTypeTest.java => arrow/DorisArrowTypeMappingTest.java} (79%) create mode 100644 fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperSerializedSchemaTest.java diff --git a/fe/fe-core/src/main/java/org/apache/doris/arrow/DorisArrowTypeMapping.java b/fe/fe-core/src/main/java/org/apache/doris/arrow/DorisArrowTypeMapping.java new file mode 100644 index 00000000000000..529bc774951f58 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/arrow/DorisArrowTypeMapping.java @@ -0,0 +1,229 @@ +// 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.doris.arrow; + +import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.thrift.TColumnDesc; + +import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; +import org.apache.arrow.vector.ZeroVector; +import org.apache.arrow.vector.complex.BaseRepeatedValueVector; +import org.apache.arrow.vector.complex.MapVector; +import org.apache.arrow.vector.types.DateUnit; +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * How FE describes a Doris column as an Arrow field. + * + *

This is the one place FE turns a Doris type into an Arrow type. Every Arrow schema FE hands a + * client -- today the {@code table_schema} column of Flight SQL {@code GetTables}, next the result + * and parameter schemas of prepared statements -- is built here, so that a client which types its + * columns from one of them reads the batches BE emits as that type. The values mirror + * {@code convert_to_arrow_type} in {@code be/src/format/arrow/arrow_row_batch.cpp}, which alone + * decides the shape of the data on the wire; the cells known to disagree with it (a literal + * {@code "UTC"} zone on TIMESTAMPTZ, {@code Null} for TIMEV2 / VARBINARY / AGG_STATE) are kept + * deliberately until the BE Arrow type layer is reworked, and are then corrected here, once, against + * a golden shared with BE (#67577). {@code DorisArrowTypeMappingTest} records every cell as it + * stands, so a change to any of them is a change to that test as well. + */ +public final class DorisArrowTypeMapping { + + private DorisArrowTypeMapping() { + } + + /** + * The Arrow type of a Doris type; {@code precision} and {@code scale} are those of the column and + * are null for a type that has none. + */ + public static ArrowType toArrowType(PrimitiveType primitiveType, Integer precision, Integer scale) { + switch (primitiveType) { + case BOOLEAN: + return new ArrowType.Bool(); + case TINYINT: + return new ArrowType.Int(8, true); + case SMALLINT: + return new ArrowType.Int(16, true); + case INT: + case IPV4: + return new ArrowType.Int(32, true); + case BIGINT: + return new ArrowType.Int(64, true); + case FLOAT: + return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); + case DOUBLE: + return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); + case LARGEINT: + case VARCHAR: + case STRING: + case CHAR: + case DATETIME: + case DATE: + case JSONB: + case IPV6: + case VARIANT: + return new ArrowType.Utf8(); + case DATEV2: + // DAY, not MILLISECOND: BE writes a DATEV2 column as arrow::Date32Type (a day number), + // so a MILLISECOND unit here describes the metadata as date64 while the data that + // follows is date32. A client that trusts this schema -- one reading through the ADBC + // Flight SQL driver does -- then types the column as a datetime and fails the read. + return new ArrowType.Date(DateUnit.DAY); + case DATETIMEV2: + if (scale > 3) { + return new ArrowType.Timestamp(TimeUnit.MICROSECOND, null); + } else if (scale > 0) { + return new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); + } else { + return new ArrowType.Timestamp(TimeUnit.SECOND, null); + } + case TIMESTAMP_NS: + return new ArrowType.Timestamp(TimeUnit.NANOSECOND, null); + case TIMESTAMPTZ: + if (scale > 3) { + return new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"); + } else if (scale > 0) { + return new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"); + } else { + return new ArrowType.Timestamp(TimeUnit.SECOND, "UTC"); + } + case DECIMAL32: + case DECIMAL64: + case DECIMAL128: + return new ArrowType.Decimal(precision, scale, 128); + case DECIMAL256: + return new ArrowType.Decimal(precision, scale, 256); + case DECIMALV2: + return new ArrowType.Decimal(27, 9, 128); + case HLL: + case BITMAP: + case QUANTILE_STATE: + return new ArrowType.Binary(); + case MAP: + return new ArrowType.Map(false); + case ARRAY: + return new ArrowType.List(); + case STRUCT: + return new ArrowType.Struct(); + default: + return new ArrowType.Null(); + } + } + + /** The Arrow type of a column as {@code describeTables} reports it. */ + public static ArrowType toArrowType(TColumnDesc desc) { + PrimitiveType primitiveType = PrimitiveType.fromThrift(desc.getColumnType()); + Integer precision = desc.isSetColumnPrecision() ? desc.getColumnPrecision() : null; + Integer scale = desc.isSetColumnScale() ? desc.getColumnScale() : null; + return toArrowType(primitiveType, precision, scale); + } + + /** + * One column of {@code dbName.tableName} as an Arrow field, with its nested types described down + * to the leaves and the Flight SQL column metadata a client's {@code ResultSetMetaData} reads. + */ + public static Field toField(String dbName, String tableName, TColumnDesc desc) { + ArrowType arrowType = toArrowType(desc); + return new Field(desc.getColumnName(), + new FieldType(desc.isIsAllowNull(), arrowType, null, + flightSqlColumnMetadata(dbName, tableName, desc)), + children(dbName, tableName, desc, arrowType)); + } + + /** + * The Arrow children of a complex column, built from the descriptor's own children. + * + *

These are not decoration. An Arrow ARRAY/MAP/STRUCT type carries its element types in its + * children and nowhere else, so a placeholder child says the column is an array OF NOTHING -- + * and BE emits the real element type in the data ({@code convert_to_arrow_type}: ListType(item), + * MapType(key, value), StructType(fields)), which leaves the schema describing one thing and the + * batch carrying another. A client that types its columns from this schema (one reading through + * the ADBC Flight SQL driver does) then rejects the column outright. + * + *

{@code describeTables} already reports the tree -- {@code Column.createChildrenColumn} names + * an array's element "item" and a map's pair "key"/"value", which is what Arrow calls them too. + * When it reports none, the old placeholders are kept rather than an empty child list: a source + * that cannot describe its nested types is no worse off than before. + */ + private static List children(String dbName, String tableName, TColumnDesc desc, + ArrowType arrowType) { + List children = desc.isSetChildren() ? desc.getChildren() : Collections.emptyList(); + switch (arrowType.getTypeID()) { + case List: + case LargeList: + case FixedSizeList: + if (children.size() != 1) { + return Collections.singletonList( + Field.notNullable(BaseRepeatedValueVector.DATA_VECTOR_NAME, + ZeroVector.INSTANCE.getField().getType())); + } + return Collections.singletonList(toField(dbName, tableName, children.get(0))); + case Map: + // Arrow spells a map as list>, with the entries struct and + // the key both non-nullable -- the descriptor's key nullability is not carried over, + // because an Arrow map with a nullable key is not a valid schema. + if (children.size() != 2) { + return Collections.singletonList( + Field.notNullable(MapVector.DATA_VECTOR_NAME, new ArrowType.List())); + } + Field key = toField(dbName, tableName, children.get(0)); + Field value = toField(dbName, tableName, children.get(1)); + Field entries = new Field(MapVector.DATA_VECTOR_NAME, + new FieldType(false, new ArrowType.Struct(), null), + Arrays.asList(new Field(key.getName(), + new FieldType(false, key.getType(), null), key.getChildren()), + value)); + return Collections.singletonList(entries); + case Struct: + if (children.isEmpty()) { + return Collections.emptyList(); + } + List structFields = new ArrayList<>(children.size()); + for (TColumnDesc child : children) { + structFields.add(toField(dbName, tableName, child)); + } + return structFields; + default: + return null; + } + } + + private static Map flightSqlColumnMetadata(final String dbName, final String tableName, + final TColumnDesc desc) { + final FlightSqlColumnMetadata.Builder columnMetadataBuilder = new FlightSqlColumnMetadata.Builder().schemaName( + dbName).tableName(tableName).typeName(PrimitiveType.fromThrift(desc.getColumnType()).toString()) + .isAutoIncrement(false).isCaseSensitive(false).isReadOnly(true).isSearchable(true); + + if (desc.isSetColumnPrecision()) { + columnMetadataBuilder.precision(desc.getColumnPrecision()); + } + if (desc.isSetColumnScale()) { + columnMetadataBuilder.scale(desc.getColumnScale()); + } + return columnMetadataBuilder.build().getMetadataMap(); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlSchemaHelper.java b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlSchemaHelper.java index 89e4042ead4cc6..e7a98cf4835e1b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlSchemaHelper.java +++ b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlSchemaHelper.java @@ -17,14 +17,13 @@ package org.apache.doris.arrowflight; +import org.apache.doris.arrow.DorisArrowTypeMapping; import org.apache.doris.catalog.Env; -import org.apache.doris.catalog.PrimitiveType; import org.apache.doris.datasource.CatalogIf; import org.apache.doris.qe.ConnectContext; import org.apache.doris.service.ExecuteEnv; import org.apache.doris.service.FrontendServiceImpl; import org.apache.doris.thrift.TColumnDef; -import org.apache.doris.thrift.TColumnDesc; import org.apache.doris.thrift.TDescribeTablesParams; import org.apache.doris.thrift.TDescribeTablesResult; import org.apache.doris.thrift.TGetDbsParams; @@ -33,23 +32,14 @@ import org.apache.doris.thrift.TListTableStatusResult; import org.apache.doris.thrift.TTableStatus; -import org.apache.arrow.flight.sql.FlightSqlColumnMetadata; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetTables; import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.ZeroVector; -import org.apache.arrow.vector.complex.BaseRepeatedValueVector; -import org.apache.arrow.vector.complex.MapVector; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; -import org.apache.arrow.vector.types.DateUnit; -import org.apache.arrow.vector.types.FloatingPointPrecision; -import org.apache.arrow.vector.types.TimeUnit; -import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; import org.apache.logging.log4j.LogManager; @@ -85,108 +75,6 @@ public FlightSqlSchemaHelper(ConnectContext context) { private static final byte[] EMPTY_SERIALIZED_SCHEMA = getSerializedSchema(Collections.emptyList()); - /** - * Convert Doris data type to an arrowType. - *

- * Ref: `convert_to_arrow_type` in be/src/util/arrow/row_batch.cpp. - * which is consistent with the type of Arrow data returned by Doris Arrow Flight Sql query. - */ - private static ArrowType getArrowType(PrimitiveType primitiveType, Integer precision, Integer scale) { - switch (primitiveType) { - case BOOLEAN: - return new ArrowType.Bool(); - case TINYINT: - return new ArrowType.Int(8, true); - case SMALLINT: - return new ArrowType.Int(16, true); - case INT: - case IPV4: - return new ArrowType.Int(32, true); - case BIGINT: - return new ArrowType.Int(64, true); - case FLOAT: - return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); - case DOUBLE: - return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); - case LARGEINT: - case VARCHAR: - case STRING: - case CHAR: - case DATETIME: - case DATE: - case JSONB: - case IPV6: - case VARIANT: - return new ArrowType.Utf8(); - case DATEV2: - // DAY, not MILLISECOND: BE writes a DATEV2 column as arrow::Date32Type (a day number), - // so a MILLISECOND unit here describes the metadata as date64 while the data that - // follows is date32. A client that trusts this schema -- one reading through the ADBC - // Flight SQL driver does -- then types the column as a datetime and fails the read. - return new ArrowType.Date(DateUnit.DAY); - case DATETIMEV2: - if (scale > 3) { - return new ArrowType.Timestamp(TimeUnit.MICROSECOND, null); - } else if (scale > 0) { - return new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); - } else { - return new ArrowType.Timestamp(TimeUnit.SECOND, null); - } - case TIMESTAMP_NS: - return new ArrowType.Timestamp(TimeUnit.NANOSECOND, null); - case TIMESTAMPTZ: - if (scale > 3) { - return new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC"); - } else if (scale > 0) { - return new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC"); - } else { - return new ArrowType.Timestamp(TimeUnit.SECOND, "UTC"); - } - case DECIMAL32: - case DECIMAL64: - case DECIMAL128: - return new ArrowType.Decimal(precision, scale, 128); - case DECIMAL256: - return new ArrowType.Decimal(precision, scale, 256); - case DECIMALV2: - return new ArrowType.Decimal(27, 9, 128); - case HLL: - case BITMAP: - case QUANTILE_STATE: - return new ArrowType.Binary(); - case MAP: - return new ArrowType.Map(false); - case ARRAY: - return new ArrowType.List(); - case STRUCT: - return new ArrowType.Struct(); - default: - return new ArrowType.Null(); - } - } - - private static ArrowType columnDescToArrowType(final TColumnDesc desc) { - PrimitiveType primitiveType = PrimitiveType.fromThrift(desc.getColumnType()); - Integer precision = desc.isSetColumnPrecision() ? desc.getColumnPrecision() : null; - Integer scale = desc.isSetColumnScale() ? desc.getColumnScale() : null; - return getArrowType(primitiveType, precision, scale); - } - - private static Map createFlightSqlColumnMetadata(final String dbName, final String tableName, - final TColumnDesc desc) { - final FlightSqlColumnMetadata.Builder columnMetadataBuilder = new FlightSqlColumnMetadata.Builder().schemaName( - dbName).tableName(tableName).typeName(PrimitiveType.fromThrift(desc.getColumnType()).toString()) - .isAutoIncrement(false).isCaseSensitive(false).isReadOnly(true).isSearchable(true); - - if (desc.isSetColumnPrecision()) { - columnMetadataBuilder.precision(desc.getColumnPrecision()); - } - if (desc.isSetColumnScale()) { - columnMetadataBuilder.scale(desc.getColumnScale()); - } - return columnMetadataBuilder.build().getMetadataMap(); - } - protected static byte[] getSerializedSchema(List fields) { if (EMPTY_SERIALIZED_SCHEMA == null && fields == null) { fields = Collections.emptyList(); @@ -287,80 +175,13 @@ private Map> buildTableToFields(String dbName, TDescribeTabl Integer tableOffset = describeTablesResult.getTablesOffset().get(tableIndex); for (; columnIndex < tableOffset; columnIndex++) { TColumnDef columnDef = describeTablesResult.getColumns().get(columnIndex); - fields.add(buildField(dbName, tableName, columnDef.getColumnDesc())); + fields.add(DorisArrowTypeMapping.toField(dbName, tableName, columnDef.getColumnDesc())); } tableToFields.put(tableName, fields); } return tableToFields; } - /** One column, with its nested types described down to the leaves. */ - private static Field buildField(String dbName, String tableName, TColumnDesc desc) { - ArrowType arrowType = columnDescToArrowType(desc); - return new Field(desc.getColumnName(), - new FieldType(desc.isIsAllowNull(), arrowType, null, - createFlightSqlColumnMetadata(dbName, tableName, desc)), - arrowChildren(dbName, tableName, desc, arrowType)); - } - - /** - * The Arrow children of a complex column, built from the descriptor's own children. - * - *

These are not decoration. An Arrow ARRAY/MAP/STRUCT type carries its element types in its - * children and nowhere else, so a placeholder child says the column is an array OF NOTHING -- - * and BE emits the real element type in the data ({@code convert_to_arrow_type}: ListType(item), - * MapType(key, value), StructType(fields)), which leaves the schema describing one thing and the - * batch carrying another. A client that types its columns from this schema (one reading through - * the ADBC Flight SQL driver does) then rejects the column outright. - * - *

{@code describeTables} already reports the tree -- {@code Column.createChildrenColumn} names - * an array's element "item" and a map's pair "key"/"value", which is what Arrow calls them too. - * When it reports none, the old placeholders are kept rather than an empty child list: a source - * that cannot describe its nested types is no worse off than before. - */ - private static List arrowChildren(String dbName, String tableName, TColumnDesc desc, - ArrowType arrowType) { - List children = desc.isSetChildren() ? desc.getChildren() : Collections.emptyList(); - switch (arrowType.getTypeID()) { - case List: - case LargeList: - case FixedSizeList: - if (children.size() != 1) { - return Collections.singletonList( - Field.notNullable(BaseRepeatedValueVector.DATA_VECTOR_NAME, - ZeroVector.INSTANCE.getField().getType())); - } - return Collections.singletonList(buildField(dbName, tableName, children.get(0))); - case Map: - // Arrow spells a map as list>, with the entries struct and - // the key both non-nullable -- the descriptor's key nullability is not carried over, - // because an Arrow map with a nullable key is not a valid schema. - if (children.size() != 2) { - return Collections.singletonList( - Field.notNullable(MapVector.DATA_VECTOR_NAME, new ArrowType.List())); - } - Field key = buildField(dbName, tableName, children.get(0)); - Field value = buildField(dbName, tableName, children.get(1)); - Field entries = new Field(MapVector.DATA_VECTOR_NAME, - new FieldType(false, new ArrowType.Struct(), null), - Arrays.asList(new Field(key.getName(), - new FieldType(false, key.getType(), null), key.getChildren()), - value)); - return Collections.singletonList(entries); - case Struct: - if (children.isEmpty()) { - return Collections.emptyList(); - } - List structFields = new ArrayList<>(children.size()); - for (TColumnDesc child : children) { - structFields.add(buildField(dbName, tableName, child)); - } - return structFields; - default: - return null; - } - } - /** * for FlightSqlProducer Schemas.GET_CATALOGS_SCHEMA */ diff --git a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperArrowTypeTest.java b/fe/fe-core/src/test/java/org/apache/doris/arrow/DorisArrowTypeMappingTest.java similarity index 79% rename from fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperArrowTypeTest.java rename to fe/fe-core/src/test/java/org/apache/doris/arrow/DorisArrowTypeMappingTest.java index 8b74a1093ef549..cebaa9ea34faf7 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperArrowTypeTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/arrow/DorisArrowTypeMappingTest.java @@ -15,56 +15,47 @@ // specific language governing permissions and limitations // under the License. -package org.apache.doris.arrowflight; +package org.apache.doris.arrow; import org.apache.doris.catalog.PrimitiveType; -import org.apache.doris.common.jmockit.Deencapsulation; import org.apache.doris.thrift.TColumnDesc; import org.apache.doris.thrift.TPrimitiveType; import org.apache.arrow.vector.complex.BaseRepeatedValueVector; import org.apache.arrow.vector.complex.MapVector; -import org.apache.arrow.vector.ipc.ReadChannel; -import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.DateUnit; import org.apache.arrow.vector.types.FloatingPointPrecision; import org.apache.arrow.vector.types.TimeUnit; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; -import java.io.ByteArrayInputStream; -import java.io.IOException; -import java.lang.reflect.Method; -import java.nio.channels.Channels; import java.util.Arrays; -import java.util.Collections; import java.util.List; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.Stream; /** - * What {@code CommandGetTables} says a column is, against what the query that follows actually - * carries. + * What FE says a column is, against what the query that follows actually carries. * *

Why these assertions matter. A Flight SQL client is entitled to type its columns from - * the schema in {@code GetTables} and then read the batches without re-deriving anything -- that is - * what the schema is for, and {@code getArrowType} is documented as mirroring - * {@code convert_to_arrow_type} in the backend. When the two disagree the client does not get a - * degraded answer, it gets a failed read: it decodes the batch as the type the metadata promised. - * So each case below pins the Arrow type BE emits, not merely "some" type. + * the schema FE gives it -- in {@code GetTables} today -- and then read the batches without + * re-deriving anything -- that is what the schema is for, and {@link DorisArrowTypeMapping} is + * documented as mirroring {@code convert_to_arrow_type} in the backend. When the two disagree the + * client does not get a degraded answer, it gets a failed read: it decodes the batch as the type the + * metadata promised. So each case below pins the Arrow type BE emits, not merely "some" type -- and + * where the mapping is known to be wrong, pins that too, so the correction is one deliberate step. * *

The descriptors are built the way {@code FrontendServiceImpl.getColumnDesc} builds them -- * a complex column carries its element types as {@link TColumnDesc} children, named "item" for an * array and "key"/"value" for a map by {@code Column.createChildrenColumn}. */ -public class FlightSqlSchemaHelperArrowTypeTest { +public class DorisArrowTypeMappingTest { private static final String DB = "test_db"; private static final String TABLE = "test_tbl"; @@ -84,16 +75,12 @@ private static TColumnDesc desc(String name, TPrimitiveType type, TColumnDesc... } private static Field buildField(TColumnDesc columnDesc) { - return Deencapsulation.invoke(FlightSqlSchemaHelper.class, "buildField", DB, TABLE, columnDesc); + return DorisArrowTypeMapping.toField(DB, TABLE, columnDesc); } - private static ArrowType arrowType(PrimitiveType type, Integer precision, Integer scale) throws Exception { - // Plain reflection rather than Deencapsulation: the descriptor of a column that has no precision - // or scale hands the mapping null for both, and that null is part of what the table pins. - Method getArrowType = FlightSqlSchemaHelper.class.getDeclaredMethod("getArrowType", - PrimitiveType.class, Integer.class, Integer.class); - getArrowType.setAccessible(true); - return (ArrowType) getArrowType.invoke(null, type, precision, scale); + /** A column that has no precision or scale hands the mapping null for both; the table pins that too. */ + private static ArrowType arrowType(PrimitiveType type, Integer precision, Integer scale) { + return DorisArrowTypeMapping.toArrowType(type, precision, scale); } private static Arguments row(PrimitiveType type, ArrowType expected) { @@ -183,7 +170,7 @@ private static Stream mapping() { @ParameterizedTest(name = "{0}({1}, {2}) is described as {3}") @MethodSource("mapping") public void everyTypeIsDescribedAsToday(PrimitiveType type, Integer precision, Integer scale, - ArrowType expected) throws Exception { + ArrowType expected) { Assertions.assertEquals(expected, arrowType(type, precision, scale)); } @@ -311,46 +298,4 @@ public void complexColumnWithoutChildrenKeepsThePlaceholder() { public void scalarColumnHasNoChildren() { Assertions.assertTrue(buildField(desc("i", TPrimitiveType.INT)).getChildren().isEmpty()); } - - /** - * The client does not see the {@link Field} objects, it sees the serialized schema in the - * {@code table_schema} column of {@code GetTables}. Asserting after a round trip through that encoding - * is what proves the element types actually reach it. - */ - @Test - public void theSerializedSchemaCarriesTheChildren() throws IOException { - byte[] serialized = FlightSqlSchemaHelper.getSerializedSchema(Collections.singletonList( - buildField(desc("a", TPrimitiveType.ARRAY, desc("item", TPrimitiveType.INT))))); - - Schema schema = MessageSerializer.deserializeSchema( - new ReadChannel(Channels.newChannel(new ByteArrayInputStream(serialized)))); - - Field array = schema.getFields().get(0); - Assertions.assertEquals(ArrowType.ArrowTypeID.List, array.getType().getTypeID()); - Assertions.assertEquals(new ArrowType.Int(32, true), array.getChildren().get(0).getType()); - } - - @Test - public void serializedSchemaDescribesScalarAndNestedTimestampNs() throws IOException { - byte[] serialized = FlightSqlSchemaHelper.getSerializedSchema(Arrays.asList( - buildField(desc("ts", TPrimitiveType.TIMESTAMP_NS)), - buildField(desc("items", TPrimitiveType.ARRAY, - desc("item", TPrimitiveType.TIMESTAMP_NS))), - buildField(desc("by_name", TPrimitiveType.MAP, - desc("key", TPrimitiveType.VARCHAR), - desc("value", TPrimitiveType.TIMESTAMP_NS))), - buildField(desc("record", TPrimitiveType.STRUCT, - desc("ts", TPrimitiveType.TIMESTAMP_NS))))); - - Schema schema = MessageSerializer.deserializeSchema( - new ReadChannel(Channels.newChannel(new ByteArrayInputStream(serialized)))); - ArrowType.Timestamp timestampNs = new ArrowType.Timestamp(TimeUnit.NANOSECOND, null); - Assertions.assertEquals(timestampNs, schema.getFields().get(0).getType()); - Assertions.assertEquals(timestampNs, - schema.getFields().get(1).getChildren().get(0).getType()); - Assertions.assertEquals(timestampNs, - schema.getFields().get(2).getChildren().get(0).getChildren().get(1).getType()); - Assertions.assertEquals(timestampNs, - schema.getFields().get(3).getChildren().get(0).getType()); - } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperSerializedSchemaTest.java b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperSerializedSchemaTest.java new file mode 100644 index 00000000000000..ebe1189ade5d82 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/FlightSqlSchemaHelperSerializedSchemaTest.java @@ -0,0 +1,102 @@ +// 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.doris.arrowflight; + +import org.apache.doris.arrow.DorisArrowTypeMapping; +import org.apache.doris.thrift.TColumnDesc; +import org.apache.doris.thrift.TPrimitiveType; + +import org.apache.arrow.vector.ipc.ReadChannel; +import org.apache.arrow.vector.ipc.message.MessageSerializer; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.nio.channels.Channels; +import java.util.Arrays; +import java.util.Collections; + +/** + * The client does not see the {@link Field} objects {@link DorisArrowTypeMapping} builds, it sees the + * serialized schema in the {@code table_schema} column of {@code GetTables}. Asserting after a round + * trip through that encoding is what proves the types, nested ones included, actually reach it. + */ +public class FlightSqlSchemaHelperSerializedSchemaTest { + + private static final String DB = "test_db"; + private static final String TABLE = "test_tbl"; + + private static TColumnDesc desc(String name, TPrimitiveType type) { + TColumnDesc columnDesc = new TColumnDesc(name, type); + columnDesc.setIsAllowNull(true); + return columnDesc; + } + + private static TColumnDesc desc(String name, TPrimitiveType type, TColumnDesc... children) { + TColumnDesc columnDesc = desc(name, type); + columnDesc.setChildren(Arrays.asList(children)); + return columnDesc; + } + + private static Field buildField(TColumnDesc columnDesc) { + return DorisArrowTypeMapping.toField(DB, TABLE, columnDesc); + } + + private static Schema deserialize(byte[] serialized) throws IOException { + return MessageSerializer.deserializeSchema( + new ReadChannel(Channels.newChannel(new ByteArrayInputStream(serialized)))); + } + + @Test + public void theSerializedSchemaCarriesTheChildren() throws IOException { + byte[] serialized = FlightSqlSchemaHelper.getSerializedSchema(Collections.singletonList( + buildField(desc("a", TPrimitiveType.ARRAY, desc("item", TPrimitiveType.INT))))); + + Field array = deserialize(serialized).getFields().get(0); + Assertions.assertEquals(ArrowType.ArrowTypeID.List, array.getType().getTypeID()); + Assertions.assertEquals(new ArrowType.Int(32, true), array.getChildren().get(0).getType()); + } + + @Test + public void serializedSchemaDescribesScalarAndNestedTimestampNs() throws IOException { + byte[] serialized = FlightSqlSchemaHelper.getSerializedSchema(Arrays.asList( + buildField(desc("ts", TPrimitiveType.TIMESTAMP_NS)), + buildField(desc("items", TPrimitiveType.ARRAY, + desc("item", TPrimitiveType.TIMESTAMP_NS))), + buildField(desc("by_name", TPrimitiveType.MAP, + desc("key", TPrimitiveType.VARCHAR), + desc("value", TPrimitiveType.TIMESTAMP_NS))), + buildField(desc("record", TPrimitiveType.STRUCT, + desc("ts", TPrimitiveType.TIMESTAMP_NS))))); + + Schema schema = deserialize(serialized); + ArrowType.Timestamp timestampNs = new ArrowType.Timestamp(TimeUnit.NANOSECOND, null); + Assertions.assertEquals(timestampNs, schema.getFields().get(0).getType()); + Assertions.assertEquals(timestampNs, + schema.getFields().get(1).getChildren().get(0).getType()); + Assertions.assertEquals(timestampNs, + schema.getFields().get(2).getChildren().get(0).getChildren().get(1).getType()); + Assertions.assertEquals(timestampNs, + schema.getFields().get(3).getChildren().get(0).getType()); + } +} From 3a90f8d7d186441b604512d69eb8c4be419f21d8 Mon Sep 17 00:00:00 2001 From: morningman Date: Mon, 21 Sep 2026 14:58:55 +0800 Subject: [PATCH 3/3] [test](arrow-flight) Pin the GetTables schema a Flight SQL client sees, one table of every type DorisArrowTypeMappingTest records the mapping as a Java table; this records what reaches a client. A raw Flight SQL client asks GetTables(include_schema) for three tables that together declare a column of every type an internal table can hold -- the scalars, DATETIME and TIMESTAMPTZ at three scales each, TIMESTAMP_NS, DECIMAL(9,2) / (18,4) / (38,10) / (76,20), JSON, VARIANT, IPV4 / IPV6, BITMAP / HLL / QUANTILE_STATE / AGG_STATE, ARRAY / MAP / STRUCT and a three-level nesting -- decodes each table_schema the way a client does and compares it, field by field and down to the leaves, with the Arrow type, the nullability and the Flight SQL column metadata as they are served today. The path it covers is the whole of it: describeTables' descriptors, DorisArrowTypeMapping, the schema serialization. The cells known to disagree with BE are pinned as they are and say so: the literal "UTC" zone on TIMESTAMPTZ and Null for AGG_STATE (TIMEV2 and VARBINARY cannot be declared on an internal table; the unit test pins those). Correcting them is one deliberate step in DorisArrowTypeMapping once the BE Arrow type layer has been reworked (#67577), and this suite is what makes a stray change to any of them visible in CI. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_get_tables_schema.groovy | 270 ++++++++++++++++++ 1 file changed, 270 insertions(+) create mode 100644 regression-test/suites/arrow_flight_sql_p0/test_get_tables_schema.groovy diff --git a/regression-test/suites/arrow_flight_sql_p0/test_get_tables_schema.groovy b/regression-test/suites/arrow_flight_sql_p0/test_get_tables_schema.groovy new file mode 100644 index 00000000000000..1bc3779539c810 --- /dev/null +++ b/regression-test/suites/arrow_flight_sql_p0/test_get_tables_schema.groovy @@ -0,0 +1,270 @@ +// 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. + +import java.nio.channels.Channels + +// The Flight SQL JDBC driver on the classpath shades Arrow Flight; its FlightSqlClient is the +// one a test can drive directly (see test_session_options). +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CloseSessionRequest +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightClient +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.Location +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.sql.FlightSqlClient +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.memory.RootAllocator +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.VarBinaryVector +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.ipc.ReadChannel +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.ipc.message.MessageSerializer +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.types.pojo.Field +import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.types.pojo.Schema + +// What a Flight SQL client is told a column is. The table_schema column of GetTables(include_schema) +// is the one statement FE makes about the Arrow type of a column, and a client that types its columns +// from it (the ADBC drivers do) reads the batches BE emits as that type: a wrong cell is not a +// degraded answer but a failed read. This pins the schema for a column of every type, nested types +// down to the leaves, exactly as it is served today -- including the cells known to disagree with +// what BE emits, which stay as they are on purpose until the BE Arrow type layer is reworked and the +// whole mapping is corrected in one step (#67577): TIMESTAMPTZ carries the literal zone "UTC" where +// BE stamps the session time zone, and AGG_STATE is described as Null where BE emits binary. +// A change to any line below is a change every GetTables client sees; make it deliberately, and make +// it in DorisArrowTypeMapping, the one place FE maps a Doris type to an Arrow type. +// +// Not in the 'arrow_flight_sql' group on purpose: `sql` stays the MySQL control connection that +// creates the tables, and GetTables is asked through a raw Flight SQL client. +suite("test_get_tables_schema") { + String host = context.config.otherConfigs.get("extArrowFlightSqlHost") + int port = context.config.otherConfigs.get("extArrowFlightSqlPort") as int + String user = context.config.otherConfigs.get("extArrowFlightSqlUser") + String password = context.config.otherConfigs.get("extArrowFlightSqlPassword") + + def db = context.dbName + def allTypes = "get_tables_schema_all_types" + def aggTypes = "get_tables_schema_agg_types" + def dec256 = "get_tables_schema_dec256" + + sql "DROP TABLE IF EXISTS ${allTypes}" + sql """ + CREATE TABLE ${allTypes} ( + k_int INT NOT NULL, + c_bool BOOLEAN, + c_tinyint TINYINT, + c_smallint SMALLINT, + c_bigint BIGINT, + c_largeint LARGEINT, + c_float FLOAT, + c_double DOUBLE, + c_decimal_9_2 DECIMAL(9, 2), + c_decimal_18_4 DECIMAL(18, 4), + c_decimal_38_10 DECIMAL(38, 10), + c_date DATE, + c_datetime_0 DATETIME(0), + c_datetime_3 DATETIME(3), + c_datetime_6 DATETIME(6), + c_timestamp_ns TIMESTAMP_NS, + c_timestamptz_0 TIMESTAMPTZ(0), + c_timestamptz_3 TIMESTAMPTZ(3), + c_timestamptz_6 TIMESTAMPTZ(6), + c_char CHAR(10), + c_varchar VARCHAR(100), + c_string STRING, + c_json JSON, + c_variant VARIANT, + c_ipv4 IPV4, + c_ipv6 IPV6, + c_array_int ARRAY, + c_array_datetime ARRAY, + c_map MAP, + c_struct STRUCT>, + c_nested ARRAY>> + ) + DUPLICATE KEY(k_int) + DISTRIBUTED BY HASH(k_int) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + sql "SET enable_agg_state = true" + sql "DROP TABLE IF EXISTS ${aggTypes}" + sql """ + CREATE TABLE ${aggTypes} ( + k_int INT NOT NULL, + c_bitmap BITMAP BITMAP_UNION, + c_hll HLL HLL_UNION, + c_quantile_state QUANTILE_STATE QUANTILE_UNION, + c_agg_state AGG_STATE GENERIC + ) + AGGREGATE KEY(k_int) + DISTRIBUTED BY HASH(k_int) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + sql "SET enable_decimal256 = true" + sql "DROP TABLE IF EXISTS ${dec256}" + sql """ + CREATE TABLE ${dec256} ( + k_int INT NOT NULL, + c_decimal_76_20 DECIMAL(76, 20) + ) + DUPLICATE KEY(k_int) + DISTRIBUTED BY HASH(k_int) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + + // One line per field, children indented under their parent: the Arrow type as Arrow Java spells + // it, "not null" where the schema says so, and in braces the Flight SQL column metadata a client's + // ResultSetMetaData reads -- the Doris type name, the precision and the scale, when present. + def describe + describe = { Field field, int depth, List out -> + def meta = field.getMetadata() + def tags = [] + if (meta.containsKey("ARROW:FLIGHT:SQL:TYPE_NAME")) { + tags << meta.get("ARROW:FLIGHT:SQL:TYPE_NAME") + } + if (meta.containsKey("ARROW:FLIGHT:SQL:PRECISION")) { + tags << "precision=" + meta.get("ARROW:FLIGHT:SQL:PRECISION") + } + if (meta.containsKey("ARROW:FLIGHT:SQL:SCALE")) { + tags << "scale=" + meta.get("ARROW:FLIGHT:SQL:SCALE") + } + def line = (" " * depth) + field.getName() + ": " + field.getType() + line += field.isNullable() ? "" : " not null" + line += tags.isEmpty() ? "" : " {" + tags.join(", ") + "}" + out << line + field.getChildren().each { describe(it, depth + 1, out) } + } + + def allocator = new RootAllocator() + def client = FlightClient.builder(allocator, Location.forGrpcInsecure(host, port)).build() + try { + def cred = client.authenticateBasicToken(user, password).get() + def flight = new FlightSqlClient(client) + + // GetTables(include_schema) for one table, its table_schema decoded the way a client decodes it. + def tableSchema = { String table -> + def info = flight.getTables("internal", db, table, null, true, cred) + def schemas = [] + info.getEndpoints().each { endpoint -> + flight.getStream(endpoint.getTicket(), cred).withCloseable { stream -> + while (stream.next()) { + def root = stream.getRoot() + def names = root.getVector("table_name") + def bytes = (VarBinaryVector) root.getVector("table_schema") + for (int i = 0; i < root.getRowCount(); i++) { + assertEquals(table, names.getObject(i).toString()) + schemas << MessageSerializer.deserializeSchema(new ReadChannel( + Channels.newChannel(new ByteArrayInputStream(bytes.get(i))))) + } + } + } + } + assertEquals(1, schemas.size(), "GetTables should describe ${db}.${table} exactly once") + return (Schema) schemas[0] + } + // The expected block is written indented for readability; the least-indented line is depth 0. + def dedent = { String text -> + def lines = text.readLines().findAll { !it.trim().isEmpty() } + int indent = lines.collect { it.length() - it.stripLeading().length() }.min() + return lines.collect { it.substring(indent) }.join("\n") + } + def check = { String table, String expected -> + Schema schema = tableSchema(table) + def lines = [] + schema.getFields().each { describe(it, 0, lines) } + assertEquals(dedent(expected), lines.join("\n"), + "the GetTables schema of ${db}.${table} changed; see DorisArrowTypeMapping") + // Every column names the table it belongs to and is read-only, as JDBC clients expect. + schema.getFields().each { field -> + def meta = field.getMetadata() + assertEquals(db, meta.get("ARROW:FLIGHT:SQL:SCHEMA_NAME"), field.getName()) + assertEquals(table, meta.get("ARROW:FLIGHT:SQL:TABLE_NAME"), field.getName()) + assertEquals("1", meta.get("ARROW:FLIGHT:SQL:IS_READ_ONLY"), field.getName()) + assertEquals("1", meta.get("ARROW:FLIGHT:SQL:IS_SEARCHABLE"), field.getName()) + assertEquals("0", meta.get("ARROW:FLIGHT:SQL:IS_AUTO_INCREMENT"), field.getName()) + assertEquals("0", meta.get("ARROW:FLIGHT:SQL:IS_CASE_SENSITIVE"), field.getName()) + } + } + + // DATETIME is a wall-clock value: a timezone-naive timestamp whose unit follows the scale. + // TIMESTAMPTZ takes the same units but the literal zone "UTC" (kept as is, see the header). + // LARGEINT is text, as BE writes it; DATE is a day number (date32), not date64. + // A map is list> with the entries struct and the key not null, + // whatever the column declares, because an Arrow map with a nullable key is not a schema. + check(allTypes, """ + k_int: Int(32, true) not null {INT, precision=10, scale=0} + c_bool: Bool {BOOLEAN, scale=0} + c_tinyint: Int(8, true) {TINYINT, precision=3, scale=0} + c_smallint: Int(16, true) {SMALLINT, precision=5, scale=0} + c_bigint: Int(64, true) {BIGINT, precision=19, scale=0} + c_largeint: Utf8 {LARGEINT, precision=39} + c_float: FloatingPoint(SINGLE) {FLOAT, precision=7, scale=7} + c_double: FloatingPoint(DOUBLE) {DOUBLE, precision=15, scale=15} + c_decimal_9_2: Decimal(9, 2, 128) {DECIMAL32, precision=9, scale=2} + c_decimal_18_4: Decimal(18, 4, 128) {DECIMAL64, precision=18, scale=4} + c_decimal_38_10: Decimal(38, 10, 128) {DECIMAL128, precision=38, scale=10} + c_date: Date(DAY) {DATEV2} + c_datetime_0: Timestamp(SECOND, null) {DATETIMEV2, precision=18, scale=0} + c_datetime_3: Timestamp(MILLISECOND, null) {DATETIMEV2, precision=18, scale=3} + c_datetime_6: Timestamp(MICROSECOND, null) {DATETIMEV2, precision=18, scale=6} + c_timestamp_ns: Timestamp(NANOSECOND, null) {TIMESTAMP_NS, precision=29, scale=9} + c_timestamptz_0: Timestamp(SECOND, UTC) {TIMESTAMPTZ, precision=18, scale=0} + c_timestamptz_3: Timestamp(MILLISECOND, UTC) {TIMESTAMPTZ, precision=18, scale=3} + c_timestamptz_6: Timestamp(MICROSECOND, UTC) {TIMESTAMPTZ, precision=18, scale=6} + c_char: Utf8 {CHAR} + c_varchar: Utf8 {VARCHAR} + c_string: Utf8 {STRING} + c_json: Utf8 {JSON} + c_variant: Utf8 {VARIANT} + c_ipv4: Int(32, true) {IPV4} + c_ipv6: Utf8 {IPV6} + c_array_int: List {ARRAY} + item: Int(32, true) {INT, precision=10, scale=0} + c_array_datetime: List {ARRAY} + item: Timestamp(MICROSECOND, null) {DATETIMEV2, precision=18, scale=6} + c_map: Map(false) {MAP} + entries: Struct not null + key: Utf8 not null + value: Int(64, true) {BIGINT, precision=19, scale=0} + c_struct: Struct {STRUCT} + f1: Int(32, true) {INT, precision=10, scale=0} + f2: Utf8 {STRING} + f3: List {ARRAY} + item: Date(DAY) {DATEV2} + c_nested: List {ARRAY} + item: Map(false) {MAP} + entries: Struct not null + key: Utf8 not null + value: List {ARRAY} + item: Decimal(10, 3, 128) {DECIMAL64, precision=10, scale=3} + """) + + // BITMAP / HLL / QUANTILE_STATE are opaque bytes. AGG_STATE is described as Null (kept as is, + // see the header). + check(aggTypes, """ + k_int: Int(32, true) not null {INT, precision=10, scale=0} + c_bitmap: Binary not null {BITMAP} + c_hll: Binary not null {HLL} + c_quantile_state: Binary not null {QUANTILE_STATE} + c_agg_state: Null not null {AGG_STATE} + """) + + check(dec256, """ + k_int: Int(32, true) not null {INT, precision=10, scale=0} + c_decimal_76_20: Decimal(76, 20, 256) {DECIMAL256, precision=76, scale=20} + """) + + // End the session the way the drivers do rather than leaving it to wait_timeout. + flight.closeSession(new CloseSessionRequest(), cred) + } finally { + client.close() + allocator.close() + } +}