<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */

#define C_LUCY_SORTFIELDWRITER
#define C_LUCY_SFWRITERELEM
#include "Lucy/Util/ToolSet.h"
#include <math.h>

#include "Lucy/Index/SortFieldWriter.h"
#include "Lucy/Index/Inverter.h"
#include "Lucy/Index/PolyReader.h"
#include "Lucy/Index/Segment.h"
#include "Lucy/Index/SegReader.h"
#include "Lucy/Index/Snapshot.h"
#include "Lucy/Index/SortCache.h"
#include "Lucy/Index/SortCache/NumericSortCache.h"
#include "Lucy/Index/SortCache/TextSortCache.h"
#include "Lucy/Index/SortReader.h"
#include "Lucy/Index/ZombieKeyedHash.h"
#include "Lucy/Plan/FieldType.h"
#include "Lucy/Plan/Schema.h"
#include "Lucy/Store/Folder.h"
#include "Lucy/Store/InStream.h"
#include "Lucy/Store/OutStream.h"
#include "Clownfish/Util/Memory.h"
#include "Lucy/Util/MemoryPool.h"
#include "Clownfish/Util/SortUtils.h"

// Prepare to read back a run.
static void
S_flip_run(SortFieldWriter *run, size_t sub_thresh, InStream *ord_in,
           InStream *ix_in, InStream *dat_in);

// Write out a sort cache.  Returns the number of unique values in the sort
// cache.
static int32_t
S_write_files(SortFieldWriter *self, OutStream *ord_out, OutStream *ix_out,
              OutStream *dat_out);

// Create an element for the sort pool.  Both the `value` and the object
// allocation itself will come from the MemoryPool, so the the element will be
// deallocated via MemPool_Release_All().
static SFWriterElem*
S_SFWriterElem_create(MemoryPool *mem_pool, Obj *value, int32_t doc_id);

SortFieldWriter*
SortFieldWriter_new(Schema *schema, Snapshot *snapshot, Segment *segment,
                    PolyReader *polyreader, String *field,
                    MemoryPool *memory_pool, size_t mem_thresh,
                    OutStream *temp_ord_out, OutStream *temp_ix_out,
                    OutStream *temp_dat_out) {
    SortFieldWriter *self
        = (SortFieldWriter*)VTable_Make_Obj(SORTFIELDWRITER);
    return SortFieldWriter_init(self, schema, snapshot, segment, polyreader,
                                field, memory_pool, mem_thresh, temp_ord_out,
                                temp_ix_out, temp_dat_out);
}

SortFieldWriter*
SortFieldWriter_init(SortFieldWriter *self, Schema *schema,
                     Snapshot *snapshot, Segment *segment,
                     PolyReader *polyreader, String *field,
                     MemoryPool *memory_pool, size_t mem_thresh,
                     OutStream *temp_ord_out, OutStream *temp_ix_out,
                     OutStream *temp_dat_out) {
    // Init.
    SortEx_init((SortExternal*)self);
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    ivars->null_ord        = -1;
    ivars->count           = 0;
    ivars->ord_start       = 0;
    ivars->ord_end         = 0;
    ivars->ix_start        = 0;
    ivars->ix_end          = 0;
    ivars->dat_start       = 0;
    ivars->dat_end         = 0;
    ivars->run_cardinality = -1;
    ivars->run_max         = -1;
    ivars->sort_cache      = NULL;
    ivars->doc_map         = NULL;
    ivars->sorted_ids      = NULL;
    ivars->run_ord         = 0;
    ivars->run_tick        = 0;
    ivars->ord_width       = 0;
    ivars->last_val        = NULL;

    // Assign.
    ivars->field        = Str_Clone(field);
    ivars->schema       = (Schema*)INCREF(schema);
    ivars->snapshot     = (Snapshot*)INCREF(snapshot);
    ivars->segment      = (Segment*)INCREF(segment);
    ivars->polyreader   = (PolyReader*)INCREF(polyreader);
    ivars->mem_pool     = (MemoryPool*)INCREF(memory_pool);
    ivars->temp_ord_out = (OutStream*)INCREF(temp_ord_out);
    ivars->temp_ix_out  = (OutStream*)INCREF(temp_ix_out);
    ivars->temp_dat_out = (OutStream*)INCREF(temp_dat_out);
    ivars->mem_thresh   = mem_thresh;

    // Derive.
    ivars->field_num = Seg_Field_Num(segment, field);
    FieldType *type = (FieldType*)CERTIFY(
                          Schema_Fetch_Type(ivars->schema, field), FIELDTYPE);
    ivars->type    = (FieldType*)INCREF(type);
    ivars->prim_id = FType_Primitive_ID(type);
    if (ivars->prim_id == FType_TEXT || ivars->prim_id == FType_BLOB) {
        ivars->var_width = true;
    }
    else {
        ivars->var_width = false;
    }
    ivars->uniq_vals = (Hash*)ZKHash_new(memory_pool, ivars->prim_id);

    return self;
}

void
SortFieldWriter_Clear_Buffer_IMP(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    if (ivars->uniq_vals) {
        if (ivars->last_val) {
            Obj *clone = Obj_Clone(ivars->last_val);
            DECREF(ivars->last_val);
            ivars->last_val = clone;
        }
        Hash_Clear(ivars->uniq_vals);
    }
    SortFieldWriter_Clear_Buffer_t super_clear_buffer
        = SUPER_METHOD_PTR(SORTFIELDWRITER, LUCY_SortFieldWriter_Clear_Buffer);
    super_clear_buffer(self);
    // Note that we have not called MemPool_Release_All() on our memory pool.
    // This is because the pool is shared amongst multiple SortFieldWriters
    // which belong to a parent SortWriter; it is the responsibility of the
    // parent SortWriter to release the memory pool once **all** of its child
    // SortFieldWriters have cleared their buffers.
}

void
SortFieldWriter_Destroy_IMP(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    DECREF(ivars->uniq_vals);
    ivars->uniq_vals = NULL;
    DECREF(ivars->field);
    DECREF(ivars->schema);
    DECREF(ivars->snapshot);
    DECREF(ivars->segment);
    DECREF(ivars->polyreader);
    DECREF(ivars->type);
    DECREF(ivars->mem_pool);
    DECREF(ivars->temp_ord_out);
    DECREF(ivars->temp_ix_out);
    DECREF(ivars->temp_dat_out);
    DECREF(ivars->ord_in);
    DECREF(ivars->ix_in);
    DECREF(ivars->dat_in);
    DECREF(ivars->sort_cache);
    DECREF(ivars->doc_map);
    FREEMEM(ivars->sorted_ids);
    SUPER_DESTROY(self, SORTFIELDWRITER);
}

int32_t
SortFieldWriter_Get_Null_Ord_IMP(SortFieldWriter *self) {
    return SortFieldWriter_IVARS(self)->null_ord;
}

int32_t
SortFieldWriter_Get_Ord_Width_IMP(SortFieldWriter *self) {
    return SortFieldWriter_IVARS(self)->ord_width;
}

static Obj*
S_find_unique_value(Hash *uniq_vals, Obj *val) {
    int32_t  hash_sum  = Obj_Hash_Sum(val);
    Obj     *uniq_val  = Hash_Find_Key(uniq_vals, val, hash_sum);
    if (!uniq_val) {
        Hash_Store(uniq_vals, val, (Obj*)CFISH_TRUE);
        uniq_val = Hash_Find_Key(uniq_vals, val, hash_sum);
    }
    return uniq_val;
}

void
SortFieldWriter_Add_IMP(SortFieldWriter *self, int32_t doc_id, Obj *value) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);

    // Uniq-ify the value, and record it for this document.
    Obj *copy = S_find_unique_value(ivars->uniq_vals, value);
    SFWriterElem *elem = S_SFWriterElem_create(ivars->mem_pool, copy, doc_id);
    SortFieldWriter_Feed(self, (Obj*)elem);
    ivars->count++;
}

void
SortFieldWriter_Add_Segment_IMP(SortFieldWriter *self, SegReader *reader,
                                I32Array *doc_map, SortCache *sort_cache) {
    if (!sort_cache) { return; }
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    SortFieldWriter *run
        = SortFieldWriter_new(ivars->schema, ivars->snapshot, ivars->segment,
                              ivars->polyreader, ivars->field, ivars->mem_pool,
                              ivars->mem_thresh, NULL, NULL, NULL);
    SortFieldWriterIVARS *const run_ivars = SortFieldWriter_IVARS(run);
    run_ivars->sort_cache = (SortCache*)INCREF(sort_cache);
    run_ivars->doc_map    = (I32Array*)INCREF(doc_map);
    run_ivars->run_max    = SegReader_Doc_Max(reader);
    run_ivars->run_cardinality = SortCache_Get_Cardinality(sort_cache);
    run_ivars->null_ord   = SortCache_Get_Null_Ord(sort_cache);
    run_ivars->run_tick   = 1;
    SortFieldWriter_Add_Run(self, (SortExternal*)run);
}

static int32_t
S_calc_width(int32_t cardinality) {
    if (cardinality <= 0x00000002)      { return 1; }
    else if (cardinality <= 0x00000004) { return 2; }
    else if (cardinality <= 0x0000000F) { return 4; }
    else if (cardinality <= 0x000000FF) { return 8; }
    else if (cardinality <= 0x0000FFFF) { return 16; }
    else                                { return 32; }
}

static void
S_write_ord(void *ords, int32_t width, int32_t doc_id, int32_t ord) {
    switch (width) {
        case 1:
            if (ord) { NumUtil_u1set(ords, doc_id); }
            else     { NumUtil_u1clear(ords, doc_id); }
            break;
        case 2:
            NumUtil_u2set(ords, doc_id, ord);
            break;
        case 4:
            NumUtil_u4set(ords, doc_id, ord);
            break;
        case 8: {
                uint8_t *ints = (uint8_t*)ords;
                ints[doc_id] = ord;
            }
            break;
        case 16: {
                uint8_t *bytes = (uint8_t*)ords;
                bytes += doc_id * sizeof(uint16_t);
                NumUtil_encode_bigend_u16(ord, &bytes);
            }
            break;
        case 32: {
                uint8_t *bytes = (uint8_t*)ords;
                bytes += doc_id * sizeof(uint32_t);
                NumUtil_encode_bigend_u32(ord, &bytes);
            }
            break;
        default:
            THROW(ERR, "Invalid width: %i32", width);
    }
}

static void
S_write_val(Obj *val, int8_t prim_id, OutStream *ix_out, OutStream *dat_out,
            int64_t dat_start) {
    if (val) {
        switch (prim_id & FType_PRIMITIVE_ID_MASK) {
            case FType_TEXT: {
                    String *string = (String*)val;
                    int64_t dat_pos = OutStream_Tell(dat_out) - dat_start;
                    OutStream_Write_I64(ix_out, dat_pos);
                    OutStream_Write_Bytes(dat_out, Str_Get_Ptr8(string),
                                          Str_Get_Size(string));
                    break;
                }
            case FType_BLOB: {
                    ByteBuf *byte_buf = (ByteBuf*)val;
                    int64_t dat_pos = OutStream_Tell(dat_out) - dat_start;
                    OutStream_Write_I64(ix_out, dat_pos);
                    OutStream_Write_Bytes(dat_out, BB_Get_Buf(byte_buf),
                                          BB_Get_Size(byte_buf));
                    break;
                }
            case FType_INT32: {
                    Integer32 *i32 = (Integer32*)val;
                    OutStream_Write_I32(dat_out, Int32_Get_Value(i32));
                    break;
                }
            case FType_INT64: {
                    Integer64 *i64 = (Integer64*)val;
                    OutStream_Write_I64(dat_out, Int64_Get_Value(i64));
                    break;
                }
            case FType_FLOAT64: {
                    Float64 *float64 = (Float64*)val;
                    OutStream_Write_F64(dat_out, Float64_Get_Value(float64));
                    break;
                }
            case FType_FLOAT32: {
                    Float32 *float32 = (Float32*)val;
                    OutStream_Write_F32(dat_out, Float32_Get_Value(float32));
                    break;
                }
            default:
                THROW(ERR, "Unrecognized primitive id: %i32", (int32_t)prim_id);
        }
    }
    else {
        switch (prim_id & FType_PRIMITIVE_ID_MASK) {
            case FType_TEXT:
            case FType_BLOB: {
                    int64_t dat_pos = OutStream_Tell(dat_out) - dat_start;
                    OutStream_Write_I64(ix_out, dat_pos);
                }
                break;
            case FType_INT32:
                OutStream_Write_I32(dat_out, 0);
                break;
            case FType_INT64:
                OutStream_Write_I64(dat_out, 0);
                break;
            case FType_FLOAT64:
                OutStream_Write_F64(dat_out, 0.0);
                break;
            case FType_FLOAT32:
                OutStream_Write_F32(dat_out, 0.0f);
                break;
            default:
                THROW(ERR, "Unrecognized primitive id: %i32", (int32_t)prim_id);
        }
    }
}

int
SortFieldWriter_Compare_IMP(SortFieldWriter *self, void *va, void *vb) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    SFWriterElemIVARS *a = SFWriterElem_IVARS(*(SFWriterElem**)va);
    SFWriterElemIVARS *b = SFWriterElem_IVARS(*(SFWriterElem**)vb);
    int32_t comparison
        = FType_null_back_compare_values(ivars->type, a->value, b->value);
    if (comparison == 0) { comparison = b->doc_id - a->doc_id; }
    return comparison;
}

static int
S_compare_doc_ids_by_ord_rev(void *context, const void *va, const void *vb) {
    SortCache *sort_cache = (SortCache*)context;
    int32_t a = *(int32_t*)va;
    int32_t b = *(int32_t*)vb;
    int32_t ord_a = SortCache_Ordinal(sort_cache, a);
    int32_t ord_b = SortCache_Ordinal(sort_cache, b);
    return ord_a - ord_b;
}

static void
S_lazy_init_sorted_ids(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    if (!ivars->sorted_ids) {
        ivars->sorted_ids
            = (int32_t*)MALLOCATE((ivars->run_max + 1) * sizeof(int32_t));
        for (int32_t i = 0, max = ivars->run_max; i <= max; i++) {
            ivars->sorted_ids[i] = i;
        }
        Sort_quicksort(ivars->sorted_ids + 1, ivars->run_max, sizeof(int32_t),
                       S_compare_doc_ids_by_ord_rev, ivars->sort_cache);
    }
}

void
SortFieldWriter_Flush_IMP(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);

    // Don't add a run unless we have data to put in it.
    if (SortFieldWriter_Buffer_Count(self) == 0) { return; }

    OutStream *const temp_ord_out = ivars->temp_ord_out;
    OutStream *const temp_ix_out  = ivars->temp_ix_out;
    OutStream *const temp_dat_out = ivars->temp_dat_out;

    SortFieldWriter_Sort_Buffer(self);
    SortFieldWriter *run
        = SortFieldWriter_new(ivars->schema, ivars->snapshot, ivars->segment,
                              ivars->polyreader, ivars->field, ivars->mem_pool,
                              ivars->mem_thresh, NULL, NULL, NULL);
    SortFieldWriterIVARS *const run_ivars = SortFieldWriter_IVARS(run);

    // Record stream starts and align.
    run_ivars->ord_start = OutStream_Align(temp_ord_out, sizeof(int64_t));
    if (ivars->var_width) {
        run_ivars->ix_start  = OutStream_Align(temp_ix_out, sizeof(int64_t));
    }
    run_ivars->dat_start = OutStream_Align(temp_dat_out, sizeof(int64_t));

    // Have the run borrow the array of elems.
    run_ivars->buffer   = ivars->buffer;
    run_ivars->buf_max  = ivars->buf_max;
    run_ivars->buf_tick = ivars->buf_tick;
    run_ivars->buf_cap  = ivars->buf_cap;

    // Write files, record stats.
    run_ivars->run_max = (int32_t)Seg_Get_Count(ivars->segment);
    run_ivars->run_cardinality = S_write_files(run, temp_ord_out, temp_ix_out,
                                               temp_dat_out);

    // Reclaim the buffer from the run and empty it.
    run_ivars->buffer    = NULL;
    run_ivars->buf_max   = 0;
    run_ivars->buf_tick  = 0;
    run_ivars->buf_cap   = 0;
    ivars->buf_tick = ivars->buf_max;
    SortFieldWriter_Clear_Buffer(self);

    // Record stream ends.
    run_ivars->ord_end = OutStream_Tell(temp_ord_out);
    if (ivars->var_width) {
        run_ivars->ix_end  = OutStream_Tell(temp_ix_out);
    }
    run_ivars->dat_end = OutStream_Tell(temp_dat_out);

    // Add the run to the array.
    SortFieldWriter_Add_Run(self, (SortExternal*)run);
}

uint32_t
SortFieldWriter_Refill_IMP(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    if (!ivars->sort_cache) { return 0; }

    // Sanity check, then reset the buffer and prepare to start loading items.
    uint32_t buf_count = SortFieldWriter_Buffer_Count(self);
    if (buf_count) {
        THROW(ERR, "Refill called but buffer contains %u32 items",
              buf_count);
    }
    SortFieldWriter_Clear_Buffer(self);
    MemPool_Release_All(ivars->mem_pool);
    S_lazy_init_sorted_ids(self);

    const int32_t    null_ord   = ivars->null_ord;
    Hash *const      uniq_vals  = ivars->uniq_vals;
    I32Array *const  doc_map    = ivars->doc_map;
    SortCache *const sort_cache = ivars->sort_cache;

    while (ivars->run_ord < ivars->run_cardinality
           && MemPool_Get_Consumed(ivars->mem_pool) < ivars->mem_thresh
          ) {
        Obj *val = SortCache_Value(sort_cache, ivars->run_ord);
        if (val) {
            Hash_Store(uniq_vals, val, (Obj*)CFISH_TRUE);
            DECREF(val);
            break;
        }
        ivars->run_ord++;
    }
    uint32_t count = 0;
    while (ivars->run_tick <= ivars->run_max) {
        int32_t raw_doc_id = ivars->sorted_ids[ivars->run_tick];
        int32_t ord = SortCache_Ordinal(sort_cache, raw_doc_id);
        if (ord != null_ord) {
            int32_t remapped = doc_map
                               ? I32Arr_Get(doc_map, raw_doc_id)
                               : raw_doc_id;
            if (remapped) {
                Obj *val = SortCache_Value(sort_cache, ord);
                SortFieldWriter_Add(self, remapped, val);
                count++;
                DECREF(val);
            }
        }
        else if (ord > ivars->run_ord) {
            break;
        }
        ivars->run_tick++;
    }
    ivars->run_ord++;
    SortFieldWriter_Sort_Buffer(self);

    if (ivars->run_ord >= ivars->run_cardinality) {
        DECREF(ivars->sort_cache);
        ivars->sort_cache = NULL;
    }

    return count;
}

void
SortFieldWriter_Flip_IMP(SortFieldWriter *self) {<commit_msg>Don't sort documents twice in SortFieldWriter#Refill<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */

#define C_LUCY_SORTFIELDWRITER
#define C_LUCY_SFWRITERELEM
#include "Lucy/Util/ToolSet.h"
#include <math.h>

#include "Lucy/Index/SortFieldWriter.h"
#include "Lucy/Index/Inverter.h"
#include "Lucy/Index/PolyReader.h"
#include "Lucy/Index/Segment.h"
#include "Lucy/Index/SegReader.h"
#include "Lucy/Index/Snapshot.h"
#include "Lucy/Index/SortCache.h"
#include "Lucy/Index/SortCache/NumericSortCache.h"
#include "Lucy/Index/SortCache/TextSortCache.h"
#include "Lucy/Index/SortReader.h"
#include "Lucy/Index/ZombieKeyedHash.h"
#include "Lucy/Plan/FieldType.h"
#include "Lucy/Plan/Schema.h"
#include "Lucy/Store/Folder.h"
#include "Lucy/Store/InStream.h"
#include "Lucy/Store/OutStream.h"
#include "Clownfish/Util/Memory.h"
#include "Lucy/Util/MemoryPool.h"
#include "Clownfish/Util/SortUtils.h"

// Prepare to read back a run.
static void
S_flip_run(SortFieldWriter *run, size_t sub_thresh, InStream *ord_in,
           InStream *ix_in, InStream *dat_in);

// Write out a sort cache.  Returns the number of unique values in the sort
// cache.
static int32_t
S_write_files(SortFieldWriter *self, OutStream *ord_out, OutStream *ix_out,
              OutStream *dat_out);

// Create an element for the sort pool.  Both the `value` and the object
// allocation itself will come from the MemoryPool, so the the element will be
// deallocated via MemPool_Release_All().
static SFWriterElem*
S_SFWriterElem_create(MemoryPool *mem_pool, Obj *value, int32_t doc_id);

SortFieldWriter*
SortFieldWriter_new(Schema *schema, Snapshot *snapshot, Segment *segment,
                    PolyReader *polyreader, String *field,
                    MemoryPool *memory_pool, size_t mem_thresh,
                    OutStream *temp_ord_out, OutStream *temp_ix_out,
                    OutStream *temp_dat_out) {
    SortFieldWriter *self
        = (SortFieldWriter*)VTable_Make_Obj(SORTFIELDWRITER);
    return SortFieldWriter_init(self, schema, snapshot, segment, polyreader,
                                field, memory_pool, mem_thresh, temp_ord_out,
                                temp_ix_out, temp_dat_out);
}

SortFieldWriter*
SortFieldWriter_init(SortFieldWriter *self, Schema *schema,
                     Snapshot *snapshot, Segment *segment,
                     PolyReader *polyreader, String *field,
                     MemoryPool *memory_pool, size_t mem_thresh,
                     OutStream *temp_ord_out, OutStream *temp_ix_out,
                     OutStream *temp_dat_out) {
    // Init.
    SortEx_init((SortExternal*)self);
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    ivars->null_ord        = -1;
    ivars->count           = 0;
    ivars->ord_start       = 0;
    ivars->ord_end         = 0;
    ivars->ix_start        = 0;
    ivars->ix_end          = 0;
    ivars->dat_start       = 0;
    ivars->dat_end         = 0;
    ivars->run_cardinality = -1;
    ivars->run_max         = -1;
    ivars->sort_cache      = NULL;
    ivars->doc_map         = NULL;
    ivars->sorted_ids      = NULL;
    ivars->run_ord         = 0;
    ivars->run_tick        = 0;
    ivars->ord_width       = 0;
    ivars->last_val        = NULL;

    // Assign.
    ivars->field        = Str_Clone(field);
    ivars->schema       = (Schema*)INCREF(schema);
    ivars->snapshot     = (Snapshot*)INCREF(snapshot);
    ivars->segment      = (Segment*)INCREF(segment);
    ivars->polyreader   = (PolyReader*)INCREF(polyreader);
    ivars->mem_pool     = (MemoryPool*)INCREF(memory_pool);
    ivars->temp_ord_out = (OutStream*)INCREF(temp_ord_out);
    ivars->temp_ix_out  = (OutStream*)INCREF(temp_ix_out);
    ivars->temp_dat_out = (OutStream*)INCREF(temp_dat_out);
    ivars->mem_thresh   = mem_thresh;

    // Derive.
    ivars->field_num = Seg_Field_Num(segment, field);
    FieldType *type = (FieldType*)CERTIFY(
                          Schema_Fetch_Type(ivars->schema, field), FIELDTYPE);
    ivars->type    = (FieldType*)INCREF(type);
    ivars->prim_id = FType_Primitive_ID(type);
    if (ivars->prim_id == FType_TEXT || ivars->prim_id == FType_BLOB) {
        ivars->var_width = true;
    }
    else {
        ivars->var_width = false;
    }
    ivars->uniq_vals = (Hash*)ZKHash_new(memory_pool, ivars->prim_id);

    return self;
}

void
SortFieldWriter_Clear_Buffer_IMP(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    if (ivars->uniq_vals) {
        if (ivars->last_val) {
            Obj *clone = Obj_Clone(ivars->last_val);
            DECREF(ivars->last_val);
            ivars->last_val = clone;
        }
        Hash_Clear(ivars->uniq_vals);
    }
    SortFieldWriter_Clear_Buffer_t super_clear_buffer
        = SUPER_METHOD_PTR(SORTFIELDWRITER, LUCY_SortFieldWriter_Clear_Buffer);
    super_clear_buffer(self);
    // Note that we have not called MemPool_Release_All() on our memory pool.
    // This is because the pool is shared amongst multiple SortFieldWriters
    // which belong to a parent SortWriter; it is the responsibility of the
    // parent SortWriter to release the memory pool once **all** of its child
    // SortFieldWriters have cleared their buffers.
}

void
SortFieldWriter_Destroy_IMP(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    DECREF(ivars->uniq_vals);
    ivars->uniq_vals = NULL;
    DECREF(ivars->field);
    DECREF(ivars->schema);
    DECREF(ivars->snapshot);
    DECREF(ivars->segment);
    DECREF(ivars->polyreader);
    DECREF(ivars->type);
    DECREF(ivars->mem_pool);
    DECREF(ivars->temp_ord_out);
    DECREF(ivars->temp_ix_out);
    DECREF(ivars->temp_dat_out);
    DECREF(ivars->ord_in);
    DECREF(ivars->ix_in);
    DECREF(ivars->dat_in);
    DECREF(ivars->sort_cache);
    DECREF(ivars->doc_map);
    FREEMEM(ivars->sorted_ids);
    SUPER_DESTROY(self, SORTFIELDWRITER);
}

int32_t
SortFieldWriter_Get_Null_Ord_IMP(SortFieldWriter *self) {
    return SortFieldWriter_IVARS(self)->null_ord;
}

int32_t
SortFieldWriter_Get_Ord_Width_IMP(SortFieldWriter *self) {
    return SortFieldWriter_IVARS(self)->ord_width;
}

static Obj*
S_find_unique_value(Hash *uniq_vals, Obj *val) {
    int32_t  hash_sum  = Obj_Hash_Sum(val);
    Obj     *uniq_val  = Hash_Find_Key(uniq_vals, val, hash_sum);
    if (!uniq_val) {
        Hash_Store(uniq_vals, val, (Obj*)CFISH_TRUE);
        uniq_val = Hash_Find_Key(uniq_vals, val, hash_sum);
    }
    return uniq_val;
}

void
SortFieldWriter_Add_IMP(SortFieldWriter *self, int32_t doc_id, Obj *value) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);

    // Uniq-ify the value, and record it for this document.
    Obj *copy = S_find_unique_value(ivars->uniq_vals, value);
    SFWriterElem *elem = S_SFWriterElem_create(ivars->mem_pool, copy, doc_id);
    SortFieldWriter_Feed(self, (Obj*)elem);
    ivars->count++;
}

void
SortFieldWriter_Add_Segment_IMP(SortFieldWriter *self, SegReader *reader,
                                I32Array *doc_map, SortCache *sort_cache) {
    if (!sort_cache) { return; }
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    SortFieldWriter *run
        = SortFieldWriter_new(ivars->schema, ivars->snapshot, ivars->segment,
                              ivars->polyreader, ivars->field, ivars->mem_pool,
                              ivars->mem_thresh, NULL, NULL, NULL);
    SortFieldWriterIVARS *const run_ivars = SortFieldWriter_IVARS(run);
    run_ivars->sort_cache = (SortCache*)INCREF(sort_cache);
    run_ivars->doc_map    = (I32Array*)INCREF(doc_map);
    run_ivars->run_max    = SegReader_Doc_Max(reader);
    run_ivars->run_cardinality = SortCache_Get_Cardinality(sort_cache);
    run_ivars->null_ord   = SortCache_Get_Null_Ord(sort_cache);
    run_ivars->run_tick   = 1;
    SortFieldWriter_Add_Run(self, (SortExternal*)run);
}

static int32_t
S_calc_width(int32_t cardinality) {
    if (cardinality <= 0x00000002)      { return 1; }
    else if (cardinality <= 0x00000004) { return 2; }
    else if (cardinality <= 0x0000000F) { return 4; }
    else if (cardinality <= 0x000000FF) { return 8; }
    else if (cardinality <= 0x0000FFFF) { return 16; }
    else                                { return 32; }
}

static void
S_write_ord(void *ords, int32_t width, int32_t doc_id, int32_t ord) {
    switch (width) {
        case 1:
            if (ord) { NumUtil_u1set(ords, doc_id); }
            else     { NumUtil_u1clear(ords, doc_id); }
            break;
        case 2:
            NumUtil_u2set(ords, doc_id, ord);
            break;
        case 4:
            NumUtil_u4set(ords, doc_id, ord);
            break;
        case 8: {
                uint8_t *ints = (uint8_t*)ords;
                ints[doc_id] = ord;
            }
            break;
        case 16: {
                uint8_t *bytes = (uint8_t*)ords;
                bytes += doc_id * sizeof(uint16_t);
                NumUtil_encode_bigend_u16(ord, &bytes);
            }
            break;
        case 32: {
                uint8_t *bytes = (uint8_t*)ords;
                bytes += doc_id * sizeof(uint32_t);
                NumUtil_encode_bigend_u32(ord, &bytes);
            }
            break;
        default:
            THROW(ERR, "Invalid width: %i32", width);
    }
}

static void
S_write_val(Obj *val, int8_t prim_id, OutStream *ix_out, OutStream *dat_out,
            int64_t dat_start) {
    if (val) {
        switch (prim_id & FType_PRIMITIVE_ID_MASK) {
            case FType_TEXT: {
                    String *string = (String*)val;
                    int64_t dat_pos = OutStream_Tell(dat_out) - dat_start;
                    OutStream_Write_I64(ix_out, dat_pos);
                    OutStream_Write_Bytes(dat_out, Str_Get_Ptr8(string),
                                          Str_Get_Size(string));
                    break;
                }
            case FType_BLOB: {
                    ByteBuf *byte_buf = (ByteBuf*)val;
                    int64_t dat_pos = OutStream_Tell(dat_out) - dat_start;
                    OutStream_Write_I64(ix_out, dat_pos);
                    OutStream_Write_Bytes(dat_out, BB_Get_Buf(byte_buf),
                                          BB_Get_Size(byte_buf));
                    break;
                }
            case FType_INT32: {
                    Integer32 *i32 = (Integer32*)val;
                    OutStream_Write_I32(dat_out, Int32_Get_Value(i32));
                    break;
                }
            case FType_INT64: {
                    Integer64 *i64 = (Integer64*)val;
                    OutStream_Write_I64(dat_out, Int64_Get_Value(i64));
                    break;
                }
            case FType_FLOAT64: {
                    Float64 *float64 = (Float64*)val;
                    OutStream_Write_F64(dat_out, Float64_Get_Value(float64));
                    break;
                }
            case FType_FLOAT32: {
                    Float32 *float32 = (Float32*)val;
                    OutStream_Write_F32(dat_out, Float32_Get_Value(float32));
                    break;
                }
            default:
                THROW(ERR, "Unrecognized primitive id: %i32", (int32_t)prim_id);
        }
    }
    else {
        switch (prim_id & FType_PRIMITIVE_ID_MASK) {
            case FType_TEXT:
            case FType_BLOB: {
                    int64_t dat_pos = OutStream_Tell(dat_out) - dat_start;
                    OutStream_Write_I64(ix_out, dat_pos);
                }
                break;
            case FType_INT32:
                OutStream_Write_I32(dat_out, 0);
                break;
            case FType_INT64:
                OutStream_Write_I64(dat_out, 0);
                break;
            case FType_FLOAT64:
                OutStream_Write_F64(dat_out, 0.0);
                break;
            case FType_FLOAT32:
                OutStream_Write_F32(dat_out, 0.0f);
                break;
            default:
                THROW(ERR, "Unrecognized primitive id: %i32", (int32_t)prim_id);
        }
    }
}

int
SortFieldWriter_Compare_IMP(SortFieldWriter *self, void *va, void *vb) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    SFWriterElemIVARS *a = SFWriterElem_IVARS(*(SFWriterElem**)va);
    SFWriterElemIVARS *b = SFWriterElem_IVARS(*(SFWriterElem**)vb);
    int32_t comparison
        = FType_null_back_compare_values(ivars->type, a->value, b->value);
    if (comparison == 0) { comparison = a->doc_id - b->doc_id; }
    return comparison;
}

static int
S_compare_doc_ids_by_ord_rev(void *context, const void *va, const void *vb) {
    SortCache *sort_cache = (SortCache*)context;
    int32_t a = *(int32_t*)va;
    int32_t b = *(int32_t*)vb;
    int32_t ord_a = SortCache_Ordinal(sort_cache, a);
    int32_t ord_b = SortCache_Ordinal(sort_cache, b);
    int32_t comparison = ord_a - ord_b;
    if (comparison == 0) { comparison = a - b; }
    return comparison;
}

static void
S_lazy_init_sorted_ids(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    if (!ivars->sorted_ids) {
        ivars->sorted_ids
            = (int32_t*)MALLOCATE((ivars->run_max + 1) * sizeof(int32_t));
        for (int32_t i = 0, max = ivars->run_max; i <= max; i++) {
            ivars->sorted_ids[i] = i;
        }
        Sort_quicksort(ivars->sorted_ids + 1, ivars->run_max, sizeof(int32_t),
                       S_compare_doc_ids_by_ord_rev, ivars->sort_cache);
    }
}

void
SortFieldWriter_Flush_IMP(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);

    // Don't add a run unless we have data to put in it.
    if (SortFieldWriter_Buffer_Count(self) == 0) { return; }

    OutStream *const temp_ord_out = ivars->temp_ord_out;
    OutStream *const temp_ix_out  = ivars->temp_ix_out;
    OutStream *const temp_dat_out = ivars->temp_dat_out;

    SortFieldWriter_Sort_Buffer(self);
    SortFieldWriter *run
        = SortFieldWriter_new(ivars->schema, ivars->snapshot, ivars->segment,
                              ivars->polyreader, ivars->field, ivars->mem_pool,
                              ivars->mem_thresh, NULL, NULL, NULL);
    SortFieldWriterIVARS *const run_ivars = SortFieldWriter_IVARS(run);

    // Record stream starts and align.
    run_ivars->ord_start = OutStream_Align(temp_ord_out, sizeof(int64_t));
    if (ivars->var_width) {
        run_ivars->ix_start  = OutStream_Align(temp_ix_out, sizeof(int64_t));
    }
    run_ivars->dat_start = OutStream_Align(temp_dat_out, sizeof(int64_t));

    // Have the run borrow the array of elems.
    run_ivars->buffer   = ivars->buffer;
    run_ivars->buf_max  = ivars->buf_max;
    run_ivars->buf_tick = ivars->buf_tick;
    run_ivars->buf_cap  = ivars->buf_cap;

    // Write files, record stats.
    run_ivars->run_max = (int32_t)Seg_Get_Count(ivars->segment);
    run_ivars->run_cardinality = S_write_files(run, temp_ord_out, temp_ix_out,
                                               temp_dat_out);

    // Reclaim the buffer from the run and empty it.
    run_ivars->buffer    = NULL;
    run_ivars->buf_max   = 0;
    run_ivars->buf_tick  = 0;
    run_ivars->buf_cap   = 0;
    ivars->buf_tick = ivars->buf_max;
    SortFieldWriter_Clear_Buffer(self);

    // Record stream ends.
    run_ivars->ord_end = OutStream_Tell(temp_ord_out);
    if (ivars->var_width) {
        run_ivars->ix_end  = OutStream_Tell(temp_ix_out);
    }
    run_ivars->dat_end = OutStream_Tell(temp_dat_out);

    // Add the run to the array.
    SortFieldWriter_Add_Run(self, (SortExternal*)run);
}

uint32_t
SortFieldWriter_Refill_IMP(SortFieldWriter *self) {
    SortFieldWriterIVARS *const ivars = SortFieldWriter_IVARS(self);
    if (!ivars->sort_cache) { return 0; }

    // Sanity check, then reset the buffer and prepare to start loading items.
    uint32_t buf_count = SortFieldWriter_Buffer_Count(self);
    if (buf_count) {
        THROW(ERR, "Refill called but buffer contains %u32 items",
              buf_count);
    }
    SortFieldWriter_Clear_Buffer(self);
    MemPool_Release_All(ivars->mem_pool);
    S_lazy_init_sorted_ids(self);

    const int32_t    null_ord   = ivars->null_ord;
    Hash *const      uniq_vals  = ivars->uniq_vals;
    I32Array *const  doc_map    = ivars->doc_map;
    SortCache *const sort_cache = ivars->sort_cache;

    while (ivars->run_ord < ivars->run_cardinality
           && MemPool_Get_Consumed(ivars->mem_pool) < ivars->mem_thresh
          ) {
        Obj *val = SortCache_Value(sort_cache, ivars->run_ord);
        if (val) {
            Hash_Store(uniq_vals, val, (Obj*)CFISH_TRUE);
            DECREF(val);
            break;
        }
        ivars->run_ord++;
    }
    uint32_t count = 0;
    while (ivars->run_tick <= ivars->run_max) {
        int32_t raw_doc_id = ivars->sorted_ids[ivars->run_tick];
        int32_t ord = SortCache_Ordinal(sort_cache, raw_doc_id);
        if (ord != null_ord) {
            int32_t remapped = doc_map
                               ? I32Arr_Get(doc_map, raw_doc_id)
                               : raw_doc_id;
            if (remapped) {
                Obj *val = SortCache_Value(sort_cache, ord);
                SortFieldWriter_Add(self, remapped, val);
                count++;
                DECREF(val);
            }
        }
        else if (ord > ivars->run_ord) {
            break;
        }
        ivars->run_tick++;
    }
    ivars->run_ord++;

    if (ivars->run_ord >= ivars->run_cardinality) {
        DECREF(ivars->sort_cache);
        ivars->sort_cache = NULL;
    }

    return count;
}

void
SortFieldWriter_Flip_IMP(SortFieldWriter *self) {
{'fraction_non_alphanumeric': 0.10324995930768813, 'fraction_numerical': 0.015273181053659594, 'mean_word_length': 2.4435310602522184, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 309, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 6, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6064567', 'n_tokens_mistral': 12746, 'n_tokens_neox': 12245, 'n_words': 2385}