Add $row_id column to faker tables

trinodb · Nov 8, 2024 · 4437041 · 4437041
1 parent ce19723
commit 4437041
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 5 deletions.
diff --git a/plugin/trino-faker/src/main/java/io/trino/plugin/faker/FakerColumnHandle.java b/plugin/trino-faker/src/main/java/io/trino/plugin/faker/FakerColumnHandle.java
@@ -20,13 +20,19 @@
 import static java.util.Objects.requireNonNull;
 
 public record FakerColumnHandle(
+        Kind kind,
         int columnIndex,
         String name,
         Type type,
         double nullProbability,
         String generator)
         implements ColumnHandle
 {
+    enum Kind {
+        REGULAR,
+        SYNTHETIC
+    }
+
     public FakerColumnHandle
     {
         requireNonNull(name, "name is null");

diff --git a/plugin/trino-faker/src/main/java/io/trino/plugin/faker/FakerMetadata.java b/plugin/trino-faker/src/main/java/io/trino/plugin/faker/FakerMetadata.java
@@ -46,6 +46,7 @@
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.spi.security.TrinoPrincipal;
 import io.trino.spi.statistics.ComputedStatistics;
+import io.trino.spi.type.BigintType;
 import io.trino.spi.type.CharType;
 import io.trino.spi.type.VarbinaryType;
 import io.trino.spi.type.VarcharType;
@@ -64,6 +65,8 @@
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static io.trino.plugin.faker.FakerColumnHandle.Kind.REGULAR;
+import static io.trino.plugin.faker.FakerColumnHandle.Kind.SYNTHETIC;
 import static io.trino.spi.StandardErrorCode.INVALID_COLUMN_PROPERTY;
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static io.trino.spi.StandardErrorCode.SCHEMA_ALREADY_EXISTS;
@@ -77,6 +80,7 @@ public class FakerMetadata
         implements ConnectorMetadata
 {
     public static final String SCHEMA_NAME = "default";
+    public static final String ROW_ID_COLUMN_NAME = "$row_id";
 
     @GuardedBy("this")
     private final List<SchemaInfo> schemas = new ArrayList<>();
@@ -259,6 +263,9 @@ public synchronized void setColumnComment(ConnectorSession session, ConnectorTab
         TableInfo oldInfo = tables.get(tableName);
         List<ColumnInfo> columns = oldInfo.columns().stream()
                 .map(columnInfo -> {
+                    if (ROW_ID_COLUMN_NAME.equals(columnInfo.handle().name())) {
+                        throw new IllegalArgumentException(String.format("Cannot set comment for %s column", ROW_ID_COLUMN_NAME));
+                    }
                     if (columnInfo.handle().equals(column)) {
                         return columnInfo.withComment(comment);
                     }
@@ -292,7 +299,8 @@ public synchronized FakerOutputTableHandle beginCreateTable(ConnectorSession ses
         double tableNullProbability = (double) tableMetadata.getProperties().getOrDefault(TableInfo.NULL_PROBABILITY_PROPERTY, schemaNullProbability);
 
         ImmutableList.Builder<ColumnInfo> columns = ImmutableList.builder();
-        for (int columnId = 0; columnId < tableMetadata.getColumns().size(); columnId++) {
+        int columnId = 0;
+        for (; columnId < tableMetadata.getColumns().size(); columnId++) {
             ColumnMetadata column = tableMetadata.getColumns().get(columnId);
             double nullProbability = 0;
             if (column.isNullable()) {
@@ -304,6 +312,7 @@ public synchronized FakerOutputTableHandle beginCreateTable(ConnectorSession ses
             }
             columns.add(new ColumnInfo(
                     new FakerColumnHandle(
+                            REGULAR,
                             columnId,
                             column.getName(),
                             column.getType(),
@@ -312,6 +321,21 @@ public synchronized FakerOutputTableHandle beginCreateTable(ConnectorSession ses
                     column));
         }
 
+        columns.add(new ColumnInfo(
+                new FakerColumnHandle(
+                        SYNTHETIC,
+                        columnId,
+                        ROW_ID_COLUMN_NAME,
+                        BigintType.BIGINT,
+                        0,
+                        ""),
+                ColumnMetadata.builder()
+                        .setName(ROW_ID_COLUMN_NAME)
+                        .setType(BigintType.BIGINT)
+                        .setHidden(true)
+                        .setNullable(false)
+                        .build()));
+
         tables.put(tableName, new TableInfo(
                 columns.build(),
                 tableMetadata.getProperties(),

diff --git a/plugin/trino-faker/src/main/java/io/trino/plugin/faker/FakerPageSource.java b/plugin/trino-faker/src/main/java/io/trino/plugin/faker/FakerPageSource.java
@@ -50,6 +50,7 @@
 import java.util.Random;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static io.trino.plugin.faker.FakerMetadata.ROW_ID_COLUMN_NAME;
 import static io.trino.spi.StandardErrorCode.INVALID_ROW_FILTER;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
@@ -119,7 +120,13 @@ class FakerPageSource
 
     private boolean closed;
 
-    FakerPageSource(Faker faker, Random random, List<FakerColumnHandle> columns, TupleDomain<ColumnHandle> constraint, long limit)
+    FakerPageSource(
+            Faker faker,
+            Random random,
+            List<FakerColumnHandle> columns,
+            TupleDomain<ColumnHandle> constraint,
+            long offset,
+            long limit)
     {
         this.faker = requireNonNull(faker, "faker is null");
         this.random = requireNonNull(random, "random is null");
@@ -132,13 +139,27 @@ class FakerPageSource
 
         this.generators = columns
                 .stream()
-                .map(column -> getGenerator(column, constraint))
+                .map(column -> getGenerator(column, constraint, offset))
                 .collect(toImmutableList());
         this.pageBuilder = new PageBuilder(types);
     }
 
-    private Generator getGenerator(FakerColumnHandle column, TupleDomain<ColumnHandle> constraint)
+    private Generator getGenerator(
+            FakerColumnHandle column,
+            TupleDomain<ColumnHandle> constraint,
+            long offset)
     {
+        if (ROW_ID_COLUMN_NAME.equals(column.name())) {
+            return new Generator() {
+                long currentRowId = offset;
+                @Override
+                public void accept(BlockBuilder blockBuilder)
+                {
+                    BIGINT.writeLong(blockBuilder, currentRowId++);
+                }
+            };
+        }
+
         return constraintedValueGenerator(
                 column,
                 constraint.getDomains().get().getOrDefault(column, Domain.all(column.type())));

diff --git a/plugin/trino-faker/src/main/java/io/trino/plugin/faker/FakerPageSourceProvider.java b/plugin/trino-faker/src/main/java/io/trino/plugin/faker/FakerPageSourceProvider.java
@@ -66,7 +66,7 @@ public ConnectorPageSource createPageSource(
         FakerTableHandle fakerTable = (FakerTableHandle) table;
         FakerSplit fakerSplit = (FakerSplit) split;
         Random random = random(fakerSplit.splitNumber());
-        return new FakerPageSource(new Faker(random), random, handles, fakerTable.constraint(), fakerSplit.rowsCount());
+        return new FakerPageSource(new Faker(random), random, handles, fakerTable.constraint(), fakerSplit.rowsOffset(), fakerSplit.rowsCount());
     }
 
     private Random random(long index)