root/platform/mysql/top-n-sort/top-n-sort.c

#include <mysql.h>
#include <pthread.h>
#include <stdlib.h>
#include <string.h>

#ifndef MAX_ITEMS
# define MAX_ITEMS 200
#endif

extern void topn_init(void) __attribute__ ((constructor));
extern void topn_fini(void) __attribute__ ((destructor));

extern my_bool topn_set_init(UDF_INIT *initid, UDF_ARGS *args, char *message);
extern void topn_set_deinit(UDF_INIT *initid);
extern void topn_set_reset(UDF_INIT *initid, UDF_ARGS *args, char *is_null,
                           char *error);
extern void topn_set_clear(UDF_INIT *initid, char *is_null, char *error);
extern void topn_set_add(UDF_INIT *initid, UDF_ARGS *args, char *is_null,
                          char *error);
extern long long topn_set(UDF_INIT *initid, UDF_ARGS *args, char *is_null,
                          char *error);
extern my_bool topn_get_init(UDF_INIT *initid, UDF_ARGS *args, char *message);
extern void topn_get_deinit(UDF_INIT *initid);
extern long long topn_get(UDF_INIT *initid, UDF_ARGS *args, char *is_null,
                          char *error);
extern my_bool topn_get_value_init(UDF_INIT *initid, UDF_ARGS *args,
                                   char *message);
extern void topn_get_value_deinit(UDF_INIT *initid);
extern long long topn_get_value(UDF_INIT *initid, UDF_ARGS *args, char *is_null,
                                char *error);

typedef struct _topn_item {
  long long id;
  long long value;
} topn_item;

typedef struct _topn_info {
  int refcnt;
  int count;
  topn_item items[MAX_ITEMS];
} topn_info;

static pthread_key_t pkey_;

void topn_init(void)
{
  pthread_key_create(&pkey_, NULL);
}

void topn_fini(void)
{
  pthread_key_delete(pkey_);
}

__inline topn_info *get_info(void)
{
  return pthread_getspecific(pkey_);
}

static topn_info *acquire_info(void)
{
  topn_info *info;
  
  if ((info = get_info()) != NULL) {
    info->refcnt++;
    return info;
  }
  
  if ((info = malloc(sizeof(*info))) == NULL) {
    return NULL;
  }
  
  info->refcnt = 1;
  info->count = 0;
  pthread_setspecific(pkey_, info);
  return info;
}

static void release_info(topn_info *info)
{
  if (--info->refcnt == 0) {
    free(info);
    pthread_setspecific(pkey_, NULL);
  }
}

#define SETUP_ARG(i, t, n)   \
  do {                       \
    args->arg_type[i] = t;   \
    args->maybe_null[i] = n; \
  } while (0)

my_bool topn_set_init(UDF_INIT *initid, UDF_ARGS *args, char *message)
{
  if (args->arg_count != 3) {
    strcpy(message, "topn_set_init(max_items,id,value): invalid arguments");
    return 1;
  }
  if ((initid->ptr = (char*)acquire_info()) == NULL) {
    strcpy(message, "no memory");
    return 1;
  }
  
  initid->maybe_null = 0;
  SETUP_ARG(0, INT_RESULT, 0);
  SETUP_ARG(1, INT_RESULT, 0);
  SETUP_ARG(2, INT_RESULT, 0);
  return 0;
}

void topn_set_deinit(UDF_INIT *initid)
{
  release_info((topn_info*)initid->ptr);
}

void topn_set_reset(UDF_INIT *initid, UDF_ARGS *args, char *is_null,
                    char *error)
{
  topn_set_clear(initid, is_null, error);
  topn_set_add(initid, args, is_null, error);
}

void topn_set_clear(UDF_INIT *initid, char *is_null, char *error)
{
  topn_info *info = (topn_info*)initid->ptr;
  info->count = 0;
}

void topn_set_add(UDF_INIT *initid, UDF_ARGS *args, char *is_null, char *error)
{
  topn_info *info = (topn_info*)initid->ptr;
  int max_items = (int)*(long long*)args->args[0];
  long long id = *(long long*)args->args[1],
    value = *(long long*)args->args[2];
  int keep, i;
  
  for (keep = info->count - 1; keep >= 0; keep--) {
    if (info->items[keep].value < value
        || (info->items[keep].value == value && info->items[keep].id < id)) {
      break;
    }
  }
  if (keep != max_items - 1) {
    if (info->count < max_items) {
      info->count++;
    }
    for (i = info->count - 2; i > keep; i--) {
      memcpy(info->items + i + 1, info->items + i, sizeof(info->items[0]));
    }
    info->items[keep + 1].id = id;
    info->items[keep + 1].value = value;
  }
}

long long topn_set(UDF_INIT *initid, UDF_ARGS *args, char *is_null, char *error)
{
  topn_info *info = (topn_info*)initid->ptr;
  return info->count;
}

my_bool topn_get_init(UDF_INIT *initid, UDF_ARGS *args, char *message)
{
  if (args->arg_count != 1) {
    strcpy(message, "topn_get(index): invalid arguments");
    return 1;
  }
  if ((initid->ptr = (char*)acquire_info()) == NULL) {
    strcpy(message, "no memory");
    return 1;
  }
  
  initid->maybe_null = 0;
  SETUP_ARG(0, INT_RESULT, 0);
  return 0;
}

void topn_get_deinit(UDF_INIT *initid)
{
  release_info((topn_info*)initid->ptr);
}

long long topn_get(UDF_INIT *initid, UDF_ARGS *args, char *is_null, char *error)
{
  topn_info *info = (topn_info*)initid->ptr;
  int idx = (int)*(long long*)args->args[0];
  
  if (idx < 0) {
    idx += info->count;
  }
  return 0 <= idx && idx < info->count ? info->items[idx].id : 0;
}

my_bool topn_get_value_init(UDF_INIT *initid, UDF_ARGS *args, char *message)
{
  if (args->arg_count != 1) {
    strcpy(message, "topn_get_value(index): invalid arguments");
    return 1;
  }
  if ((initid->ptr = (char*)acquire_info()) == NULL) {
    strcpy(message, "no memory");
    return 1;
  }
  
  initid->maybe_null = 0;
  SETUP_ARG(0, INT_RESULT, 0);
  return 0;
}

void topn_get_value_deinit(UDF_INIT *initid)
{
  release_info((topn_info*)initid->ptr);
}

long long topn_get_value(UDF_INIT *initid, UDF_ARGS *args, char *is_null,
                         char *error)
{
  topn_info *info = (topn_info*)initid->ptr;
  int idx = (int)*(long long*)args->args[0];
  
  if (idx < 0) {
    idx += info->count;
  }
  return 0 <= idx && idx < info->count ? info->items[idx].value : 0;
}