diff --git a/ingester-protocol/src/main/java/io/greptime/models/ArrowHelper.java b/ingester-protocol/src/main/java/io/greptime/models/ArrowHelper.java index cbbef3f..86560f3 100644 --- a/ingester-protocol/src/main/java/io/greptime/models/ArrowHelper.java +++ b/ingester-protocol/src/main/java/io/greptime/models/ArrowHelper.java @@ -22,8 +22,10 @@ import io.greptime.v1.Common; import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.HashMap; import java.util.Iterator; import java.util.List; +import java.util.Map; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.BitVector; import org.apache.arrow.vector.DateDayVector; @@ -61,6 +63,7 @@ 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; /** @@ -68,6 +71,20 @@ */ public class ArrowHelper { + /** + * The Arrow field metadata key that carries the semantic type of a column. + * The value must be "tag", "field" or "timestamp". Recognized by the + * GreptimeDB server when auto-creating tables for bulk insert requests. + */ + public static final String SEMANTIC_TYPE_METADATA_KEY = "greptime:semantic_type"; + + /** + * The Arrow field metadata key that carries the GreptimeDB extended data type + * of a column, e.g. "Json". Recognized by the GreptimeDB server when + * auto-creating tables for bulk insert requests. + */ + public static final String DATA_TYPE_METADATA_KEY = "greptime:type"; + /** * Get the Arrow compression type from the context. * @@ -89,6 +106,12 @@ public static ArrowCompressionType getArrowCompressionType(Context ctx) { /** * Create an Arrow schema from a table schema. + *

+ * The semantic types of the columns are encoded as Arrow field metadata so that + * the GreptimeDB server can auto-create the table for bulk insert requests: + * each column carries {@code greptime:semantic_type} with value {@code tag}, + * {@code field} or {@code timestamp}, the timestamp (time index) column is + * non-nullable, and JSON columns additionally carry {@code greptime:type=Json}. * * @param tableSchema the table schema * @return the Arrow schema @@ -101,6 +124,7 @@ public static Schema createSchema(TableSchema tableSchema) { List fields = new ArrayList<>(columnCount); List columnNames = tableSchema.getColumnNames(); + List semanticTypes = tableSchema.getSemanticTypes(); List dataTypes = tableSchema.getDataTypes(); List dataTypeExtensions = tableSchema.getDataTypeExtensions(); @@ -108,8 +132,30 @@ public static Schema createSchema(TableSchema tableSchema) { String name = columnNames.get(i); ArrowType type = convertToArrowType(dataTypes.get(i), dataTypeExtensions.get(i)); - Field field = Field.nullable(name, type); - fields.add(field); + // JSON columns carry a second metadata entry (greptime:type=Json) + Map metadata = new HashMap<>(4); + boolean nullable = true; + switch (semanticTypes.get(i)) { + case TIMESTAMP: + metadata.put(SEMANTIC_TYPE_METADATA_KEY, "timestamp"); + // The time index column is never null in GreptimeDB, and the server + // requires it to be non-nullable when auto-creating the table. + nullable = false; + break; + case TAG: + metadata.put(SEMANTIC_TYPE_METADATA_KEY, "tag"); + break; + case FIELD: + metadata.put(SEMANTIC_TYPE_METADATA_KEY, "field"); + break; + default: + throw new IllegalArgumentException("Unsupported semantic type: " + semanticTypes.get(i)); + } + if (dataTypes.get(i) == Common.ColumnDataType.JSON) { + metadata.put(DATA_TYPE_METADATA_KEY, "Json"); + } + + fields.add(new Field(name, new FieldType(nullable, type, null, metadata), null)); } return new Schema(fields); diff --git a/ingester-protocol/src/test/java/io/greptime/models/ArrowHelperTest.java b/ingester-protocol/src/test/java/io/greptime/models/ArrowHelperTest.java index eaae8bc..7456e12 100644 --- a/ingester-protocol/src/test/java/io/greptime/models/ArrowHelperTest.java +++ b/ingester-protocol/src/test/java/io/greptime/models/ArrowHelperTest.java @@ -20,6 +20,7 @@ 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.Assert; import org.junit.Test; @@ -150,4 +151,36 @@ public void testCreateSchema() { new ArrowType.Binary().getTypeID(), schema.getFields().get(22).getType().getTypeID()); } + + @Test + public void testCreateSchemaWithSemanticTypeMetadata() { + TableSchema tableSchema = TableSchema.newBuilder("my_table") + .addTag("tag1", DataType.String) + .addTimestamp("ts", DataType.TimestampMillisecond) + .addField("field1", DataType.Float64) + .addField("field2", DataType.Json) + .build(); + + Schema schema = ArrowHelper.createSchema(tableSchema); + + Field tagField = schema.getFields().get(0); + Assert.assertTrue(tagField.isNullable()); + Assert.assertEquals("tag", tagField.getMetadata().get(ArrowHelper.SEMANTIC_TYPE_METADATA_KEY)); + Assert.assertFalse(tagField.getMetadata().containsKey(ArrowHelper.DATA_TYPE_METADATA_KEY)); + + Field tsField = schema.getFields().get(1); + Assert.assertFalse(tsField.isNullable()); + Assert.assertEquals("timestamp", tsField.getMetadata().get(ArrowHelper.SEMANTIC_TYPE_METADATA_KEY)); + Assert.assertFalse(tsField.getMetadata().containsKey(ArrowHelper.DATA_TYPE_METADATA_KEY)); + + Field field1 = schema.getFields().get(2); + Assert.assertTrue(field1.isNullable()); + Assert.assertEquals("field", field1.getMetadata().get(ArrowHelper.SEMANTIC_TYPE_METADATA_KEY)); + Assert.assertFalse(field1.getMetadata().containsKey(ArrowHelper.DATA_TYPE_METADATA_KEY)); + + Field jsonField = schema.getFields().get(3); + Assert.assertTrue(jsonField.isNullable()); + Assert.assertEquals("field", jsonField.getMetadata().get(ArrowHelper.SEMANTIC_TYPE_METADATA_KEY)); + Assert.assertEquals("Json", jsonField.getMetadata().get(ArrowHelper.DATA_TYPE_METADATA_KEY)); + } } diff --git a/ingester-rpc/src/main/java/io/greptime/rpc/Context.java b/ingester-rpc/src/main/java/io/greptime/rpc/Context.java index e4e74d8..f9c4cb7 100644 --- a/ingester-rpc/src/main/java/io/greptime/rpc/Context.java +++ b/ingester-rpc/src/main/java/io/greptime/rpc/Context.java @@ -28,6 +28,8 @@ @SuppressWarnings({"unchecked"}) public class Context { + private static final String AUTO_CREATE_TABLE_HINT = "auto_create_table"; + private final Map ctx = new HashMap<>(); private Compression compression = Compression.None; @@ -90,6 +92,10 @@ public Context with(String key, Object value) { /** * Adds a hint to the context. + *

+ * Repeated calls append entries rather than overwriting: calling this twice + * with the same key produces a hint string like {@code key=value1,key=value2}, + * and which value takes effect is up to server-side parsing. * * @param key the key * @param value the value @@ -109,6 +115,29 @@ public Context withHint(String key, String value) { return this; } + /** + * Sets whether this request allows the server to auto-create a missing table. + *

+ * The hint travels with any write that uses this {@link Context}: it is sent + * on bulk writes (Flight) as well as on row-based writes, where it is carried + * as a gRPC header by {@code ContextToHeadersInterceptor}. + *

+ * Note the asymmetry: {@code false} always disables auto-creation for this + * request, while {@code true} cannot override a server whose global + * auto-create-table option is disabled. + *

+ * Like {@link #withHint}, repeated calls append instead of overwriting, so + * calling this twice on a reused {@link Context} yields duplicate + * {@code auto_create_table} hints. Use a fresh {@link Context} (or + * {@link #remove(String)} on the hints key) to change the value. + * + * @param enabled whether this request allows automatic table creation + * @return this {@link Context} + */ + public Context withAutoCreateTable(boolean enabled) { + return withHint(AUTO_CREATE_TABLE_HINT, Boolean.toString(enabled)); + } + /** * Gets the hints from the context. * diff --git a/ingester-rpc/src/test/java/io/greptime/rpc/ContextTest.java b/ingester-rpc/src/test/java/io/greptime/rpc/ContextTest.java index 2113ed2..96701e4 100644 --- a/ingester-rpc/src/test/java/io/greptime/rpc/ContextTest.java +++ b/ingester-rpc/src/test/java/io/greptime/rpc/ContextTest.java @@ -51,6 +51,13 @@ public void withHintTest() { Assert.assertEquals("key=value,key2=value2", context.getHints()); } + @Test + public void withAutoCreateTableTest() { + Context context = Context.newDefault().withAutoCreateTable(false); + + Assert.assertEquals("auto_create_table=false", context.getHints()); + } + @Test public void getShouldReturnNullForNonExistingKeyTest() { Context context = Context.newDefault();