在StarRocks中使用UDWF计算用户留存

create table user_activity (uid int, date string)
    engine = olap duplicate key (`uid`)
    distributed by hash(`uid`)
    properties ("replication_num" = "1");

def main():
    data = []
    today = datetime.date.today()
    for user_id in range(3):
        for d in range(10):
            if random.random() < 0.5:
                date = today - datetime.timedelta(days=d)
                data.append((date.strftime('%Y-%m-%d'), user_id))

    with open('user_activity.csv', 'w', newline='') as csvfile:
        fieldnames = ['date', 'user_id']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)

        # writer.writeheader()
        for row in data:
            writer.writerow({'date': row[0], 'user_id': row[1]})

main()

curl --location-trusted -u root: -H "label:124" \
    -H "column_separator:," \
    -H "columns: date,uid" \
    -T user_activity.csv -XPUT \
    http://127.0.0.1:41001/api/zya/user_activity/_stream_load

with Rte as (select uid, date as date , RetentionEncode(date)  over (partition by uid order by date desc) as encoded from user_activity)
     select date, count(1) as cohort,
          sum(bitand(bit_shift_right(encoded, 1), 1)) as day1,
          sum(bitand(bit_shift_right(encoded, 2), 1)) as day2 from Rte
          group by date order by date;

drop global function RetentionEncode(string);

CREATE GLOBAL AGGREGATE FUNCTION RetentionEncode(string)
RETURNS bigint
properties
(
    "analytic" = "true",
    "symbol" = "com.starrocks.udf.retentionEncode",
    "type" = "StarrocksJar",
    "file" = "http://127.0.0.1:41006/data/sr-udf-1.0-SNAPSHOT-jar-with-dependencies.jar"
);

package com.starrocks.udf;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.TreeSet;
import java.util.concurrent.TimeUnit;

public class retentionEncode {
    public static class State {
        TreeSet<Long> buffer = new TreeSet<>();

        Long encoded = 0L;

        @Override
        public String toString() {
            return String.format("State(encoded = 0x%x, buffer = %d)", encoded, buffer.size());
        }

        public int serializeLength() {
            return 4;
        }
    }

    public State create() {
        return new State();
    }

    public void destroy(State state) {

    }

    public void update(State state, String val) {
    }

    public void serialize(State state, java.nio.ByteBuffer buff) {
    }

    public void merge(State state, java.nio.ByteBuffer buffer) {
    }

    public Long finalize(State state) {
        System.out.println("finalize: " + state.toString());
        return state.encoded;
    }

    public void reset(State state) {
        System.out.println("reset: " + state.toString());
        state.buffer.clear();
        state.encoded = 0L;
    }

    private static Date stringToDate(String val) {
        SimpleDateFormat dateFormat = new SimpleDateFormat("yyyy-MM-dd");
        try {
            Date date = dateFormat.parse(val);
            return date;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    private void updateEncoded(State state, Date val) {
        state.encoded = 0L;
        if (val == null) {
            return;
        }

        long now = val.getTime();
        state.buffer.add(now);

        final int maxDays = 8;

        // pop out date >= maxDays
        while (state.buffer.size() > 0) {
            long last = state.buffer.last();
            long diffInMs = last - now;
            long diffInDays = TimeUnit.DAYS.convert(diffInMs, TimeUnit.MILLISECONDS);
            if (diffInDays >= maxDays) {
                state.buffer.pollLast();
            } else {
                break;
            }
        }

        for (long t : state.buffer) {
            long diffInMs = t - now;
            long diffInDays = TimeUnit.DAYS.convert(diffInMs, TimeUnit.MILLISECONDS);
            assert diffInDays < maxDays;
            state.encoded |= (1L << diffInDays);
        }
    }

    public void windowUpdate(State state, int peer_group_start, int peer_group_end, int frame_start, int frame_end,
                             String[] dates) {
        Date val = stringToDate(dates[frame_start]);
        updateEncoded(state, val);
    }
}