[WIP] AUC Plugin

lib/trino-array/TrueraTrinoPlugin.java

@@ -0,0 +1,2 @@
+public class TrueraTrinoPlugin {
+}

plugin/trino-truera/README.md

@@ -0,0 +1,19 @@
+# TruEra Trino Extensions
+This module is a custom Trino plugin for TruEra. Currently it just contains one new function (ROC-AUC).
+
+## How do I add to the plugin?
+To get started, read the Trino [dev guide](https://trino.io/docs/current/develop/spi-overview.html#)
+
+## How do I test the package?
+1. Compile using `mvn clean install -Dair.check.skip-all=true -DskipTests`.
+2. Look at the `target` folder which should now have a packaged ZIP like: `trino-truera-389.zip`
+3. Unzip the package into the trino plugins directory: `unzip trino-truera-389.zip -d ~/external_dependencies/trino-server/plugin`
+4. Restart Trino (`./service.sh stop trino && ./service.sh start trino`)
+
+You can test the "roc_auc" function with a command like:
+```sql
+SELECT roc_auc(__truera_prediction__, CAST(__truera_label__ as boolean)) 
+FROM "iceberg"."tablestore"."a83db5d335ab494590cd7ada132707ad_predictions_probits_score" as pred
+JOIN "iceberg"."tablestore"."a83db5d335ab494590cd7ada132707ad_label" 
+AS label ON pred.__truera_id__ = label.__truera_id__;
+```

plugin/trino-truera/pom.xml

@@ -0,0 +1,106 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <parent>
+        <artifactId>trino-root</artifactId>
+        <groupId>io.trino</groupId>
+        <version>406</version>
+        <relativePath>../../pom.xml</relativePath>
+    </parent>
+
+    <artifactId>trino-truera</artifactId>
+    <description>Trino Truera Extensions</description>
+    <packaging>trino-plugin</packaging>
+
+    <properties>
+        <air.main.basedir>${project.parent.basedir}</air.main.basedir>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-array</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-spi</artifactId>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>slice</artifactId>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>log</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.openjdk.jol</groupId>
+            <artifactId>jol-core</artifactId>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.testng</groupId>
+            <artifactId>testng</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-main</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-main</artifactId>
+            <type>test-jar</type>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-main</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-testing</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-hive</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-iceberg</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-hive-hadoop2</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-main</artifactId>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <pluginManagement>
+            <plugins>
+                <plugin>
+                    <groupId>com.mycila</groupId>
+                    <artifactId>license-maven-plugin</artifactId>
+                    <configuration>
+                        <excludes combine.children="append">
+                            <exclude>src/main/**</exclude>
+                        </excludes>
+                    </configuration>
+                </plugin>
+            </plugins>
+        </pluginManagement>
+    </build>
+
+</project>

plugin/trino-truera/src/main/java/io/trino/plugin/truera/AreaUnderRocCurveAlgorithm.java

@@ -0,0 +1,46 @@
+package io.trino.plugin.truera;
+
+import java.util.Comparator;
+import io.airlift.log.Logger;
+import java.util.stream.IntStream;
+
+public class AreaUnderRocCurveAlgorithm {
+    private static final Logger log = Logger.get(AreaUnderRocCurveAlgorithm.class);
+    public static double computeRocAuc(boolean[] labels, double[] scores) {
+        log.info("cow");
+        log.info("compute", labels.toString(), scores.toString());
+        int[] sortedIndices = IntStream.range(0, scores.length).boxed().sorted(
+                Comparator.comparing(i -> scores[i], Comparator.reverseOrder())
+        ).mapToInt(i->i).toArray();
+
+        int currTruePositives = 0, currFalsePositives = 0;
+        double auc = 0.;
+
+        int i = 0;
+        while (i < sortedIndices.length) {
+            int prevTruePositives = currTruePositives;
+            int prevFalsePositives = currFalsePositives;
+            double currentScore = scores[sortedIndices[i]];
+            while (i < sortedIndices.length && currentScore == scores[sortedIndices[i]]) {
+                if (labels[sortedIndices[i]]) {
+                    currTruePositives++;
+                } else {
+                    currFalsePositives++;
+                }
+                ++i;
+            }
+            auc += trapezoidIntegrate(prevFalsePositives, currFalsePositives, prevTruePositives, currTruePositives);
+        }
+
+        // If labels only contain one class, AUC is undefined
+        if (currTruePositives == 0 || currFalsePositives == 0) {
+            return Double.NaN;
+        }
+
+        return auc / (currTruePositives * currFalsePositives);
+    }
+
+    private static double trapezoidIntegrate(double x1, double x2, double y1, double y2) {
+        return (y1 + y2) * Math.abs(x2 - x1) / 2; // (base1 + base2) * height / 2
+    }
+}

plugin/trino-truera/src/main/java/io/trino/plugin/truera/TrueraTrinoPlugin.java

@@ -0,0 +1,30 @@
+/*
+ * Licensed 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 io.trino.plugin.truera;
+
+import io.trino.spi.Plugin;
+import io.trino.plugin.truera.aggregation.ROCAUCAggregation;
+
+import java.util.Collections;
+import java.util.Set;
+
+public class TrueraTrinoPlugin
+        implements Plugin
+{
+    @Override
+    public Set<Class<?>> getFunctions()
+    {
+        return Collections.singleton(ROCAUCAggregation.class);
+    }
+}

plugin/trino-truera/src/main/java/io/trino/plugin/truera/aggregation/AUCBlockBuilderStatus.java

@@ -0,0 +1,5 @@
+package io.trino.plugin.truera.aggregation;
+
+import io.trino.spi.block.BlockBuilder;
+import io.trino.spi.block.BlockBuilderStatus;
+import io.trino.spi.block.PageBuilderStatus;

plugin/trino-truera/src/main/java/io/trino/plugin/truera/aggregation/GroupedRocAucCurve.java

@@ -0,0 +1,150 @@
+package io.trino.plugin.truera.aggregation;
+
+import io.airlift.log.Logger;
+import java.util.ArrayList;
+import java.util.List;
+
+
+import io.trino.array.BooleanBigArray;
+import io.trino.array.IntBigArray;
+import io.trino.array.LongBigArray;
+import io.trino.array.DoubleBigArray;
+
+import io.trino.spi.block.Block;
+import io.trino.spi.block.BlockBuilder;
+import io.trino.spi.type.BooleanType;
+import io.trino.spi.type.DoubleType;
+import org.openjdk.jol.info.ClassLayout;
+
+import static io.trino.plugin.truera.AreaUnderRocCurveAlgorithm.computeRocAuc;
+import static java.util.Objects.requireNonNull;
+
+public class GroupedRocAucCurve {
+    private static final Logger log = Logger.get(GroupedRocAucCurve.class);
+
+    private static final long INSTANCE_SIZE = ClassLayout.parseClass(GroupedRocAucCurve.class).instanceSize();
+    private static final int NULL = -1;
+
+    // one entry per group
+    // each entry is the index of the first elements of the group in the labels/scores/nextLinks arrays
+    private final LongBigArray headIndices;
+
+    // one entry per double/boolean pair
+    private final BooleanBigArray labels;
+    private final DoubleBigArray scores;
+
+    // the value in nextLinks contains the index of the next value in the chain
+    // a value NULL (-1) indicates it is the last value for the group
+    private final IntBigArray nextLinks;
+
+    // the index of the next free element in the labels/scores/nextLinks arrays
+    // this is needed to be able to know where to continue adding elements when after the arrays are resized
+    private int nextFreeIndex;
+
+    private long currentGroupId = -1;
+
+    public GroupedRocAucCurve() {
+        this.headIndices = new LongBigArray(NULL);
+        this.labels = new BooleanBigArray();
+        this.scores = new DoubleBigArray();
+        this.nextLinks = new IntBigArray(NULL);
+        this.nextFreeIndex = 0;
+    }
+
+    public GroupedRocAucCurve(long groupId, Block serialized) {
+        this();
+        this.currentGroupId = groupId;
+
+        requireNonNull(serialized, "serialized block is null");
+        for (int i = 0; i < serialized.getPositionCount(); i++) {
+            Block entryBlock = serialized.getObject(i, Block.class);
+            add(entryBlock, entryBlock, 0, 1);
+        }
+    }
+
+    public void serialize(BlockBuilder out) {
+        if (isCurrentGroupEmpty()) {
+            out.appendNull();
+            return;
+        }
+
+        // retrieve scores + labels
+        List<Boolean> labelList = new ArrayList<>();
+        List<Double> scoreList = new ArrayList<>();
+
+        int currentIndex = (int) headIndices.get(currentGroupId);
+        while (currentIndex != NULL) {
+            labelList.add(labels.get(currentIndex));
+            scoreList.add(scores.get(currentIndex));
+            currentIndex = nextLinks.get(currentIndex);
+        }
+
+        // convert lists to primitive arrays
+        boolean[] labels = new boolean[labelList.size()];
+        for (int i = 0; i < labels.length; i++) {
+            labels[i] = labelList.get(i);
+        }
+        double[] scores = scoreList.stream().mapToDouble(Double::doubleValue).toArray();
+        log.info("cow2");
+        log.info("compute", labels.toString(), scores.toString());
+
+        // compute + return
+        double auc = computeRocAuc(labels, scores);
+        if (Double.isNaN(auc)) {
+            out.appendNull();
+        } else {
+            DoubleType.DOUBLE.writeDouble(out, auc);
+        }
+    }
+
+    public long estimatedInMemorySize() {
+        return INSTANCE_SIZE + labels.sizeOf() + scores.sizeOf() + nextLinks.sizeOf() + headIndices.sizeOf();
+    }
+
+    public GroupedRocAucCurve setGroupId(long groupId) {
+        this.currentGroupId = groupId;
+        return this;
+    }
+
+    public long getGroupId() {
+        return this.currentGroupId;
+    }
+
+    public void add(Block labelsBlock, Block scoresBlock, int labelPosition, int scorePosition) {
+        ensureCapacity(currentGroupId + 1);
+
+        labels.set(nextFreeIndex, BooleanType.BOOLEAN.getBoolean(labelsBlock, labelPosition));
+        scores.set(nextFreeIndex, DoubleType.DOUBLE.getDouble(scoresBlock, scorePosition));
+        nextLinks.set(nextFreeIndex, (int) headIndices.get(currentGroupId));
+        nextFreeIndex++;
+    }
+
+    public void ensureCapacity(long numberOfGroups) {
+        headIndices.ensureCapacity(numberOfGroups);
+        int numberOfValues = nextFreeIndex + 1;
+        labels.ensureCapacity(numberOfValues);
+        scores.ensureCapacity(numberOfValues);
+        nextLinks.ensureCapacity(numberOfValues);
+    }
+
+    public void addAll(GroupedRocAucCurve other) {
+        other.readAll(this);
+    }
+
+    public void readAll(GroupedRocAucCurve to) {
+        int currentIndex = (int) headIndices.get(currentGroupId);
+        while (currentIndex != NULL) {
+            BlockBuilder labelBlockBuilder = BooleanType.BOOLEAN.createBlockBuilder(null, 0);
+            BooleanType.BOOLEAN.writeBoolean(labelBlockBuilder, labels.get(currentIndex));
+            BlockBuilder scoreBlockBuilder = DoubleType.DOUBLE.createBlockBuilder(null, 0);
+            DoubleType.DOUBLE.writeDouble(scoreBlockBuilder, scores.get(currentIndex));
+
+            to.add(labelBlockBuilder.build(), scoreBlockBuilder.build(), 0, 0);
+            currentIndex = nextLinks.get(currentIndex);
+        }
+    }
+
+    public boolean isCurrentGroupEmpty() {
+        return headIndices.get(currentGroupId) == NULL;
+    }
+}