Multithreaded indexing ?

[lhoestq]

Is full text search indexing actually multithreaded ? On my side on a 64 core machine it only uses 2, and it tales 1min+

(I'm indexing about 5GB of text)

I see some CTEs in the source code, so this can be related to issue: duckdb/duckdb#11565
EDIT: actually it seems tokenization is not multithreaded, so probably unrelated to CTE

[lhoestq]

I could confirm the tokenization step is not multithreaded.

I ended up using python to start multiple threads to start parallel tokenizations, and add a step to group the tokenized texts together.

It would be awesome to enable multithreading for such operations !

For those interested, here is what I ended up using:

from tqdm.contrib.concurrent import thread_map

num_jobs = min(16, max(1, count // 4))
batch_size = 1 + count // num_jobs
commands = [
    (
            "ATTACH '%database%' as db (READ_ONLY);"
            "USE db;"
            f"ATTACH '{tmp_dir}/tmp_{rank}_{i}.duckdb' as tmp_{rank}_{i};"
            f"""
            CREATE TABLE tmp_{rank}_{i}.tokenized AS (
                SELECT unnest(%fts_schema%.tokenize(fts_ii."{column}")) AS w,
                    {rank * batch_size} + row_number() OVER () - 1 AS docid,
                    {i} AS fieldid
                FROM (
                    SELECT * FROM %input_table% LIMIT {batch_size} OFFSET {rank * batch_size}
                ) AS fts_ii
            );
            CHECKPOINT;
            """
    )
    for rank in range(num_jobs)
    for i, column in enumerate(columns)
]

def _parallel_sql(command: str) -> None:
    with duckdb.connect(":memory:") as rank_con:
        _sql(rank_con, command)

thread_map(_parallel_sql, commands, desc="Tokenize")

Full code:

Details

import tempfile
from textwrap import dedent

import duckdb
from tqdm.contrib.concurrent import thread_map

INSTALL_AND_LOAD_EXTENSION_COMMAND = "INSTALL 'fts'; LOAD 'fts';"
SET_EXTENSIONS_DIRECTORY_COMMAND = "SET extension_directory='{directory}';"

def create_index(
    database: str,
    input_table: str,
    columns: list[str],
    stemmer: str,
    input_id: str,
    fts_schema: str,
    extensions_directory: Optional[str] = None,
) -> None:
    """See https://github.com/duckdb/duckdb-fts/issues/7"""
    placeholders: dict[str, str] = {
        "database": database,
        "input_table": input_table,
        "stemmer": stemmer,
        "input_id": input_id,
        "fts_schema": fts_schema,
    }

    def _sql(con: duckdb.DuckDBPyConnection, query: str) -> duckdb.DuckDBPyRelation:
        query = dedent(query)
        for key, value in placeholders.items():
            query = query.replace(f"%{key}%", value)
        out = con.sql(query)
        return out

    with tempfile.TemporaryDirectory(suffix=".duckdb") as tmp_dir:
        with duckdb.connect(":memory:") as con:
            # configure duckdb extensions
            if extensions_directory is not None:
                con.execute(SET_EXTENSIONS_DIRECTORY_COMMAND.format(directory=extensions_directory))
            con.execute(INSTALL_AND_LOAD_EXTENSION_COMMAND)

            # init
            _sql(con, "ATTACH '%database%' as db;")
            _sql(con, "USE db;")

            # check input_table and get number of rows
            _count = _sql(con, "SELECT count(*) FROM %input_table%;").fetchone()
            if _count and isinstance(_count[0], int):
                count = _count[0]
            else:
                raise RuntimeError(f"failed to count rows from {input_table}")

            # create fts schema
            _sql(con, "DROP SCHEMA IF EXISTS %fts_schema% CASCADE;")
            _sql(con, "CREATE SCHEMA %fts_schema%;")

            # define stopwords
            _sql(con, "CREATE TABLE %fts_schema%.stopwords (sw VARCHAR);")
            _sql(
                con,
                "INSERT INTO %fts_schema%.stopwords VALUES ('a'), ('a''s'), ('able'), ('about'), ('above'), ('according'), ('accordingly'), ('across'), ('actually'), ('after'), ('afterwards'), ('again'), ('against'), ('ain''t'), ('all'), ('allow'), ('allows'), ('almost'), ('alone'), ('along'), ('already'), ('also'), ('although'), ('always'), ('am'), ('among'), ('amongst'), ('an'), ('and'), ('another'), ('any'), ('anybody'), ('anyhow'), ('anyone'), ('anything'), ('anyway'), ('anyways'), ('anywhere'), ('apart'), ('appear'), ('appreciate'), ('appropriate'), ('are'), ('aren''t'), ('around'), ('as'), ('aside'), ('ask'), ('asking'), ('associated'), ('at'), ('available'), ('away'), ('awfully'), ('b'), ('be'), ('became'), ('because'), ('become'), ('becomes'), ('becoming'), ('been'), ('before'), ('beforehand'), ('behind'), ('being'), ('believe'), ('below'), ('beside'), ('besides'), ('best'), ('better'), ('between'), ('beyond'), ('both'), ('brief'), ('but'), ('by'), ('c'), ('c''mon'), ('c''s'), ('came'), ('can'), ('can''t'), ('cannot'), ('cant'), ('cause'), ('causes'), ('certain'), ('certainly'), ('changes'), ('clearly'), ('co'), ('com'), ('come'), ('comes'), ('concerning'), ('consequently'), ('consider'), ('considering'), ('contain'), ('containing'), ('contains'), ('corresponding'), ('could'), ('couldn''t'), ('course'), ('currently'), ('d'), ('definitely'), ('described'), ('despite'), ('did'), ('didn''t'), ('different'), ('do'), ('does'), ('doesn''t'), ('doing'), ('don''t'), ('done'), ('down'), ('downwards'), ('during'), ('e'), ('each'), ('edu'), ('eg'), ('eight'), ('either'), ('else'), ('elsewhere'), ('enough'), ('entirely'), ('especially'), ('et'), ('etc'), ('even'), ('ever'), ('every'), ('everybody'), ('everyone'), ('everything'), ('everywhere'), ('ex'), ('exactly'), ('example'), ('except'), ('f'), ('far'), ('few'), ('fifth'), ('first'), ('five'), ('followed'), ('following'), ('follows'), ('for'), ('former'), ('formerly'), ('forth'), ('four'), ('from'), ('further'), ('furthermore'), ('g'), ('get'), ('gets'), ('getting'), ('given'), ('gives'), ('go'), ('goes'), ('going'), ('gone'), ('got'), ('gotten'), ('greetings'), ('h'), ('had'), ('hadn''t'), ('happens'), ('hardly'), ('has'), ('hasn''t'), ('have'), ('haven''t'), ('having'), ('he'), ('he''s'), ('hello'), ('help'), ('hence'), ('her'), ('here'), ('here''s'), ('hereafter'), ('hereby'), ('herein'), ('hereupon'), ('hers'), ('herself'), ('hi'), ('him'), ('himself'), ('his'), ('hither'), ('hopefully'), ('how'), ('howbeit'), ('however'), ('i'), ('i''d'), ('i''ll'), ('i''m'), ('i''ve'), ('ie'), ('if'), ('ignored'), ('immediate'), ('in'), ('inasmuch'), ('inc'), ('indeed'), ('indicate'), ('indicated'), ('indicates'), ('inner'), ('insofar'), ('instead'), ('into'), ('inward'), ('is'), ('isn''t'), ('it'), ('it''d'), ('it''ll'), ('it''s'), ('its'), ('itself'), ('j'), ('just'), ('k'), ('keep'), ('keeps'), ('kept'), ('know'), ('knows'), ('known'), ('l'), ('last'), ('lately'), ('later'), ('latter'), ('latterly'), ('least'), ('less'), ('lest'), ('let'), ('let''s'), ('like'), ('liked'), ('likely'), ('little'), ('look'), ('looking'), ('looks'), ('ltd'), ('m'), ('mainly'), ('many'), ('may'), ('maybe'), ('me'), ('mean'), ('meanwhile'), ('merely'), ('might'), ('more'), ('moreover'), ('most'), ('mostly'), ('much'), ('must'), ('my'), ('myself'), ('n'), ('name'), ('namely'), ('nd'), ('near'), ('nearly'), ('necessary'), ('need'), ('needs'), ('neither'), ('never'), ('nevertheless'), ('new'), ('next'), ('nine'), ('no'), ('nobody'), ('non'), ('none'), ('noone'), ('nor'), ('normally'), ('not'), ('nothing'), ('novel'), ('now'), ('nowhere'), ('o'), ('obviously'), ('of'), ('off'), ('often'), ('oh'), ('ok'), ('okay'), ('old'), ('on'), ('once'), ('one'), ('ones'), ('only'), ('onto'), ('or'), ('other'), ('others'), ('otherwise'), ('ought'), ('our'), ('ours'), ('ourselves'), ('out'), ('outside'), ('over'), ('overall'), ('own');",
            )
            _sql(
                con,
                "INSERT INTO %fts_schema%.stopwords VALUES ('p'), ('particular'), ('particularly'), ('per'), ('perhaps'), ('placed'), ('please'), ('plus'), ('possible'), ('presumably'), ('probably'), ('provides'), ('q'), ('que'), ('quite'), ('qv'), ('r'), ('rather'), ('rd'), ('re'), ('really'), ('reasonably'), ('regarding'), ('regardless'), ('regards'), ('relatively'), ('respectively'), ('right'), ('s'), ('said'), ('same'), ('saw'), ('say'), ('saying'), ('says'), ('second'), ('secondly'), ('see'), ('seeing'), ('seem'), ('seemed'), ('seeming'), ('seems'), ('seen'), ('self'), ('selves'), ('sensible'), ('sent'), ('serious'), ('seriously'), ('seven'), ('several'), ('shall'), ('she'), ('should'), ('shouldn''t'), ('since'), ('six'), ('so'), ('some'), ('somebody'), ('somehow'), ('someone'), ('something'), ('sometime'), ('sometimes'), ('somewhat'), ('somewhere'), ('soon'), ('sorry'), ('specified'), ('specify'), ('specifying'), ('still'), ('sub'), ('such'), ('sup'), ('sure'), ('t'), ('t''s'), ('take'), ('taken'), ('tell'), ('tends'), ('th'), ('than'), ('thank'), ('thanks'), ('thanx'), ('that'), ('that''s'), ('thats'), ('the'), ('their'), ('theirs'), ('them'), ('themselves'), ('then'), ('thence'), ('there'), ('there''s'), ('thereafter'), ('thereby'), ('therefore'), ('therein'), ('theres'), ('thereupon'), ('these'), ('they'), ('they''d'), ('they''ll'), ('they''re'), ('they''ve'), ('think'), ('third'), ('this'), ('thorough'), ('thoroughly'), ('those'), ('though'), ('three'), ('through'), ('throughout'), ('thru'), ('thus'), ('to'), ('together'), ('too'), ('took'), ('toward'), ('towards'), ('tried'), ('tries'), ('truly'), ('try'), ('trying'), ('twice'), ('two'), ('u'), ('un'), ('under'), ('unfortunately'), ('unless'), ('unlikely'), ('until'), ('unto'), ('up'), ('upon'), ('us'), ('use'), ('used'), ('useful'), ('uses'), ('using'), ('usually'), ('uucp'), ('v'), ('value'), ('various'), ('very'), ('via'), ('viz'), ('vs'), ('w'), ('want'), ('wants'), ('was'), ('wasn''t'), ('way'), ('we'), ('we''d'), ('we''ll'), ('we''re'), ('we''ve'), ('welcome'), ('well'), ('went'), ('were'), ('weren''t'), ('what'), ('what''s'), ('whatever'), ('when'), ('whence'), ('whenever'), ('where'), ('where''s'), ('whereafter'), ('whereas'), ('whereby'), ('wherein'), ('whereupon'), ('wherever'), ('whether'), ('which'), ('while'), ('whither'), ('who'), ('who''s'), ('whoever'), ('whole'), ('whom'), ('whose'), ('why'), ('will'), ('willing'), ('wish'), ('with'), ('within'), ('without'), ('won''t'), ('wonder'), ('would'), ('would'), ('wouldn''t'), ('x'), ('y'), ('yes'), ('yet'), ('you'), ('you''d'), ('you''ll'), ('you''re'), ('you''ve'), ('your'), ('yours'), ('yourself'), ('yourselves'), ('z'), ('zero');",
            )

            # define tokenize macro
            _sql(
                con,
                "CREATE MACRO %fts_schema%.tokenize(s) AS string_split_regex(regexp_replace(lower(strip_accents(s::VARCHAR)), '[^a-z]', ' ', 'g'), '\s+');",
            )

            # create fields table
            field_values = ", ".join(f"({i}, '{field}')" for i, field in enumerate(columns))
            _sql(
                con,
                """
                CREATE TABLE %fts_schema%.docs AS (
                    SELECT rowid AS docid,
                        "%input_id%" AS name
                    FROM %input_table%
                );
            """,
            )
            _sql(con, "CREATE TABLE %fts_schema%.fields (fieldid BIGINT, field VARCHAR);")
            _sql(
                con,
                f"INSERT INTO %fts_schema%.fields VALUES {field_values};",  # nosec - field_values is safe
            )
            _sql(con, "CHECKPOINT;")

        # tokenize in parallel (see https://github.com/duckdb/duckdb-fts/issues/7)
        num_jobs = min(16, max(1, count // 4))
        batch_size = 1 + count // num_jobs
        commands = [
            (
                (
                    SET_EXTENSIONS_DIRECTORY_COMMAND.format(directory=extensions_directory)
                    if extensions_directory is not None
                    else ""
                )
                + INSTALL_AND_LOAD_EXTENSION_COMMAND
                + (
                    "ATTACH '%database%' as db (READ_ONLY);"  # nosec - tmp_dir, batch_size, rank and i are safe
                    "USE db;"
                    f"ATTACH '{tmp_dir}/tmp_{rank}_{i}.duckdb' as tmp_{rank}_{i};"
                    f"""
                    CREATE TABLE tmp_{rank}_{i}.tokenized AS (
                        SELECT unnest(%fts_schema%.tokenize(fts_ii."{column}")) AS w,
                            {rank * batch_size} + row_number() OVER () - 1 AS docid,
                            {i} AS fieldid
                        FROM (
                            SELECT * FROM %input_table% LIMIT {batch_size} OFFSET {rank * batch_size}
                        ) AS fts_ii
                    );
                    CHECKPOINT;
                    """
                )
            )
            for rank in range(num_jobs)
            for i, column in enumerate(columns)
        ]

        def _parallel_sql(command: str) -> None:
            with duckdb.connect(":memory:") as rank_con:
                _sql(rank_con, command)

        thread_map(_parallel_sql, commands, desc="Tokenize")

        # # NON-PARALEL VERSION HERE FOR DOCUMENTATION:
        #
        # for i, column in enumerate(columns):
        #     _sql(con, f"""
        #         CREATE TABLE tmp.tokenized_{i} AS (
        #             SELECT unnest(%fts_schema%.tokenize(fts_ii."{column}")) AS w,
        #                 rowid AS docid,
        #                 {i} AS fieldid
        #             FROM %input_table% AS fts_ii
        #         )
        #     """)
        # union_fields_query = " UNION ALL ".join(f"SELECT * FROM tmp.tokenized_{i}" for i in range(len(columns)))

        with duckdb.connect(":memory:") as con:
            # configure duckdb extensions
            if extensions_directory is not None:
                con.execute(SET_EXTENSIONS_DIRECTORY_COMMAND.format(directory=extensions_directory))
            con.execute(INSTALL_AND_LOAD_EXTENSION_COMMAND)

            # init
            _sql(con, f"ATTACH '{tmp_dir}/tmp.duckdb' as tmp;")  # nosec - tmp_dir is safe
            _sql(con, "ATTACH '%database%' as db;")
            _sql(con, "USE db;")
            _sql(
                con,
                ";".join(
                    f"ATTACH '{tmp_dir}/tmp_{rank}_{i}.duckdb' as tmp_{rank}_{i} (READ_ONLY);"  # nosec - tmp_dir, rank and i are safe
                    for rank in range(num_jobs)
                    for i in range(len(columns))
                ),
            )

            # merge tokenizations
            union_fields_query = " UNION ALL ".join(
                f"SELECT * FROM tmp_{rank}_{i}.tokenized"  # nosec - rank and i are safe
                for rank in range(num_jobs)
                for i in range(len(columns))
            )
            _sql(con, f"CREATE TABLE tmp.tokenized AS {union_fields_query}")

            # step and stop
            _sql(
                con,
                """
                CREATE TABLE tmp.stemmed_stopped AS (
                    SELECT stem(t.w, '%stemmer%') AS term,
                        t.docid AS docid,
                        t.fieldid AS fieldid
                    FROM tmp.tokenized AS t
                    WHERE t.w NOT NULL
                    AND len(t.w) > 0
                    AND t.w NOT IN (SELECT sw FROM %fts_schema%.stopwords)
                )
            """,
            )

            # create terms table
            _sql(
                con,
                """
                CREATE TABLE %fts_schema%.terms AS (
                    SELECT ss.term,
                        ss.docid,
                        ss.fieldid
                    FROM tmp.stemmed_stopped AS ss
                )
            """,
            )

            # add doc lengths
            _sql(con, "ALTER TABLE %fts_schema%.docs ADD len BIGINT;")
            _sql(
                con,
                """
                UPDATE %fts_schema%.docs d
                SET len = (
                    SELECT count(term)
                    FROM %fts_schema%.terms AS t
                    WHERE t.docid = d.docid
                );
            """,
            )

            # create dictionary
            _sql(
                con,
                """
                CREATE TABLE tmp.distinct_terms AS (
                    SELECT DISTINCT term
                    FROM %fts_schema%.terms
                    ORDER BY docid, term
                )
            """,
            )
            _sql(
                con,
                """
                CREATE TABLE %fts_schema%.dict AS (
                    SELECT row_number() OVER () - 1 AS termid,
                    dt.term
                    FROM tmp.distinct_terms AS dt
                )
            """,
            )
            _sql(con, "ALTER TABLE %fts_schema%.terms ADD termid BIGINT;")
            _sql(
                con,
                """
                UPDATE %fts_schema%.terms t
                SET termid = (
                    SELECT termid
                    FROM %fts_schema%.dict d
                    WHERE t.term = d.term
                );
            """,
            )
            _sql(con, "ALTER TABLE %fts_schema%.terms DROP term;")

            # compute df
            _sql(con, "ALTER TABLE %fts_schema%.dict ADD df BIGINT;")
            _sql(
                con,
                """
                UPDATE %fts_schema%.dict d
                SET df = (
                    SELECT count(distinct docid)
                    FROM %fts_schema%.terms t
                    WHERE d.termid = t.termid
                    GROUP BY termid
                );
            """,
            )

            # compute stats
            _sql(
                con,
                """
                CREATE TABLE %fts_schema%.stats AS (
                    SELECT COUNT(docs.docid) AS num_docs,
                        SUM(docs.len) / COUNT(docs.len) AS avgdl
                    FROM %fts_schema%.docs AS docs
                );
            """,
            )

            # define match_bm25
            _sql(
                con,
                """
                CREATE MACRO %fts_schema%.match_bm25(docname, query_string, fields := NULL, k := 1.2, b := 0.75, conjunctive := false) AS (
                    WITH tokens AS (
                        SELECT DISTINCT stem(unnest(%fts_schema%.tokenize(query_string)), '%stemmer%') AS t
                    ),
                    fieldids AS (
                        SELECT fieldid
                        FROM %fts_schema%.fields
                        WHERE CASE WHEN fields IS NULL THEN 1 ELSE field IN (SELECT * FROM (SELECT UNNEST(string_split(fields, ','))) AS fsq) END
                    ),
                    qtermids AS (
                        SELECT termid
                        FROM %fts_schema%.dict AS dict,
                            tokens
                        WHERE dict.term = tokens.t
                    ),
                    qterms AS (
                        SELECT termid,
                            docid
                        FROM %fts_schema%.terms AS terms
                        WHERE CASE WHEN fields IS NULL THEN 1 ELSE fieldid IN (SELECT * FROM fieldids) END
                        AND termid IN (SELECT qtermids.termid FROM qtermids)
                    ),
                    term_tf AS (
                        SELECT termid,
                                docid,
                            COUNT(*) AS tf
                        FROM qterms
                        GROUP BY docid,
                                termid
                    ),
                    cdocs AS (
                        SELECT docid
                        FROM qterms
                        GROUP BY docid
                        HAVING CASE WHEN conjunctive THEN COUNT(DISTINCT termid) = (SELECT COUNT(*) FROM tokens) ELSE 1 END
                    ),
                    subscores AS (
                        SELECT docs.docid,
                            len,
                            term_tf.termid,
                            tf,
                            df,
                            (log(((SELECT num_docs FROM %fts_schema%.stats) - df + 0.5) / (df + 0.5) + 1) * ((tf * (k + 1)/(tf + k * (1 - b + b * (len / (SELECT avgdl FROM %fts_schema%.stats))))))) AS subscore
                        FROM term_tf,
                            cdocs,
                            %fts_schema%.docs AS docs,
                            %fts_schema%.dict AS dict
                        WHERE term_tf.docid = cdocs.docid
                        AND term_tf.docid = docs.docid
                        AND term_tf.termid = dict.termid
                    ),
                    scores AS (
                        SELECT docid,
                            sum(subscore) AS score
                        FROM subscores
                        GROUP BY docid
                    )
                    SELECT score
                    FROM scores,
                        %fts_schema%.docs AS docs
                    WHERE scores.docid = docs.docid
                    AND docs.name = docname
                );
            """,
            )
            _sql(con, "CHECKPOINT;")