Introduce new "id cache" for node tables

It is sometimes useful to know whether a way contains nodes of a certain
type, i.e. with certain tags. This commit adds a new functionality
called "id caches" that allows to find out.

When defining a node table in Lua with `define_table` the `ids` section
can now contain a field `cache = true`. If this is set the ids of all
nodes written to that table are stored in memory. Later, when processing
ways, that cache can be queried with `tablename:in_id_cache()`. Usually
this would be used to query all the member nodes of a way:
`...:in_id_cache(object.nodes)`. The `in_id_cache()` function returns
the indexes into the `object.nodes` array where the id matches. From
there it is possible to get the id of the matching node (with
`object.nodes[idx]`) or the location of that node (with
`object:as_point(idx)`. This information can be stored in a way table
which will be correctly updated in append mode. (In append mode the
in-memory id cache is populated with all the ids of the table for which
the cache was defined.)

The whole thing only works for node tables and querying only works in
the process_way() function. It can possibly be extended later to work
in other contexts.

Author: joto

flex-config/turning-circles.lua

@@ -0,0 +1,69 @@
+-- This config example file is released into the Public Domain.
+
+-- Create a table with turning circles that can be styled in sync with the
+-- highway they are on.
+
+local turning_circles = osm2pgsql.define_table({
+    name = 'turning_circles',
+    ids = { type = 'node', id_column = 'node_id', cache = true },
+    columns = {
+        { column = 'geom', type = 'point', not_null = true },
+    }
+})
+
+local highways = osm2pgsql.define_table({
+    name = 'highways',
+    ids = { type = 'way', id_column = 'way_id' },
+    columns = {
+        { column = 'htype', type = 'text', not_null = true },
+        { column = 'geom', type = 'linestring', not_null = true },
+    }
+})
+
+-- This table will contain entries for all node/way combinations where the way
+-- is tagged as "highway" and the node is tagged as "highway=turning_circle".
+-- The "htype" column contains the highway type, the "geom" the geometry of
+-- the node. This can be used, for instance, to draw the point in a style that
+-- fits with the style of the highway.
+--
+-- Note that you might have multiple entries for the same node in this table
+-- if it is in several ways. I that case you might have to decide at rendering
+-- time which of them to render.
+local highway_ends = osm2pgsql.define_table({
+    name = 'highway_ends',
+    ids = { type = 'way', id_column = 'way_id' },
+    columns = {
+        { column = 'htype', type = 'text', not_null = true },
+        { column = 'node_id', type = 'int8', not_null = true },
+        { column = 'geom', type = 'point', not_null = true },
+    }
+})
+
+function osm2pgsql.process_node(object)
+    if object.tags.highway == 'turning_circle' then
+        -- This insert will add the entry to the id cache later read with
+        -- in_id_cache().
+        turning_circles:insert({
+            geom = object:as_point(),
+        })
+    end
+end
+
+function osm2pgsql.process_way(object)
+    local t = object.tags.highway
+    if t then
+        highways:insert({
+            htype = t,
+            geom = object:as_linestring(),
+        })
+        local c = turning_circles:in_id_cache(object.nodes)
+        for _, n in ipairs(c) do
+            highway_ends:insert({
+                htype = t,
+                node_id = object.nodes[n],
+                geom = object:as_point(n),
+            })
+        end
+    end
+end
+

src/debug-output.cpp

@@ -61,6 +61,7 @@ void write_table_list_to_debug_log(std::vector<flex_table_t> const &tables)
         log_debug("  - data_tablespace={}", table.data_tablespace());
         log_debug("  - index_tablespace={}", table.index_tablespace());
         log_debug("  - cluster={}", table.cluster_by_geom());
+        log_debug("  - id_cache={}", table.with_id_cache());
         for (auto const &index : table.indexes()) {
             log_debug("  - INDEX USING {}", index.method());
             if (index.name().empty()) {

src/flex-lua-table.cpp

@@ -174,6 +174,17 @@ void setup_flex_table_id_columns(lua_State *lua_state, flex_table_t *table)
         throw fmt_error("Unknown ids type: {}.", type);
     }
 
+    bool const cache =
+        luaX_get_table_bool(lua_state, "cache", -1, "The ids", false);
+    lua_pop(lua_state, 1); // "cache"
+    if (cache) {
+        if (type == "node") {
+            table->enable_id_cache();
+        } else {
+            throw std::runtime_error{"ID cache only available for node ids."};
+        }
+    }
+
     std::string const name =
         luaX_get_table_string(lua_state, "id_column", -1, "The ids field");
     lua_pop(lua_state, 1); // "id_column"
@@ -459,6 +470,7 @@ void lua_wrapper_table_t::init(lua_State *lua_state)
     luaX_set_up_metatable(lua_state, "Table", OSM2PGSQL_TABLE_CLASS,
                           {{"__tostring", lua_trampoline_table_tostring},
                            {"insert", lua_trampoline_table_insert},
+                           {"in_id_cache", lua_trampoline_table_in_id_cache},
                            {"name", lua_trampoline_table_name},
                            {"schema", lua_trampoline_table_schema},
                            {"cluster", lua_trampoline_table_cluster},

src/flex-table.cpp

@@ -263,6 +263,10 @@ void flex_table_t::analyze(pg_conn_t const &db_connection) const
     analyze_table(db_connection, schema(), name());
 }
 
+void flex_table_t::enable_id_cache() noexcept { m_with_id_cache = true; }
+
+bool flex_table_t::with_id_cache() const noexcept { return m_with_id_cache; }
+
 namespace {
 
 void enable_check_trigger(pg_conn_t const &db_connection,

src/flex-table.hpp

@@ -215,6 +215,10 @@ class flex_table_t
 
     void analyze(pg_conn_t const &db_connection) const;
 
+    void enable_id_cache() noexcept;
+
+    bool with_id_cache() const noexcept;
+
 private:
     /// The schema this table is in
     std::string m_schema;
@@ -271,6 +275,9 @@ class flex_table_t
     /// Index should be a primary key.
     bool m_primary_key_index = false;
 
+    /// Do we want an ID cache for this table?
+    bool m_with_id_cache = false;
+
 }; // class flex_table_t
 
 class table_connection_t

src/idlist.cpp

@@ -22,6 +22,11 @@ osmid_t idlist_t::pop_id()
     return id;
 }
 
+bool idlist_t::contains(osmid_t id) const
+{
+    return std::binary_search(m_list.begin(), m_list.end(), id);
+}
+
 void idlist_t::sort_unique()
 {
     std::sort(m_list.begin(), m_list.end());

src/idlist.hpp

@@ -62,6 +62,13 @@ class idlist_t
 
     void reserve(std::size_t size) { m_list.reserve(size); }
 
+    /**
+     * Is the specified id in the list?
+     *
+     * You must have called sort_unique() before calling this.
+     */
+    bool contains(osmid_t id) const;
+
     /**
      * Remove id at the end of the list and return it.
      *

src/output-flex.cpp

@@ -95,6 +95,7 @@ TRAMPOLINE(app_as_geometrycollection, as_geometrycollection)
 } // anonymous namespace
 
 TRAMPOLINE(table_insert, insert)
+TRAMPOLINE(table_in_id_cache, in_id_cache)
 
 prepared_lua_function_t::prepared_lua_function_t(lua_State *lua_state,
                                                  calling_context context,
@@ -276,6 +277,9 @@ void flush_tables(std::vector<table_connection_t> &table_connections)
     for (auto &table : table_connections) {
         table.flush();
     }
+    for (auto &table : table_connections) {
+        table.sync();
+    }
 }
 
 void create_expire_tables(std::vector<expire_output_t> const &expire_outputs,
@@ -795,6 +799,10 @@ int output_flex_t::table_insert()
     auto const &object = check_and_get_context_object(table);
     osmid_t const id = table.map_id(object.type(), object.id());
 
+    if (table.with_id_cache()) {
+        get_id_cache(table).push_back(id);
+    }
+
     table_connection.new_line();
     auto *copy_mgr = table_connection.copy_mgr();
 
@@ -829,6 +837,56 @@ int output_flex_t::table_insert()
     return 1;
 }
 
+int output_flex_t::table_in_id_cache()
+{
+    if (m_calling_context != calling_context::process_way) {
+        throw std::runtime_error{
+            "The function in_id_cache() can only be called (directly or "
+            "indirectly) from the process_[untagged_]way() function."};
+    }
+
+    auto const num_params = lua_gettop(lua_state());
+    if (num_params != 2) {
+        throw std::runtime_error{
+            "Need two parameters: The osm2pgsql.Table and the id(s)."};
+    }
+
+    // The first parameter is the table object.
+    auto &table_connection = m_table_connections.at(
+        idx_from_param(lua_state(), OSM2PGSQL_TABLE_CLASS));
+    lua_remove(lua_state(), 1); // table
+
+    if (!table_connection.table().with_id_cache()) {
+        throw fmt_error("No ID cache on table {}.",
+                        table_connection.table().name());
+    }
+
+    std::vector<osmid_t> ids;
+    int const type = lua_type(lua_state(), 1);
+    if (type == LUA_TTABLE && luaX_is_array(lua_state())) {
+        luaX_for_each(lua_state(),
+                      [&]() { ids.push_back(lua_tointeger(lua_state(), -1)); });
+    } else {
+        throw std::runtime_error{"Second parameter must be an array of ids."};
+    }
+
+    auto const &cache = get_id_cache(table_connection.table());
+    lua_createtable(lua_state(), 0, 0);
+
+    lua_Integer n = 0;
+    lua_Integer idx = 1;
+    for (auto const id : ids) {
+        if (cache.contains(id)) {
+            lua_pushinteger(lua_state(), ++n);
+            lua_pushinteger(lua_state(), idx);
+            lua_rawset(lua_state(), -3);
+        }
+        ++idx;
+    }
+
+    return 1;
+}
+
 void output_flex_t::call_lua_function(prepared_lua_function_t func)
 {
     lua_pushvalue(lua_state(), func.index());
@@ -986,6 +1044,28 @@ void output_flex_t::after_nodes()
     }
 
     flush_tables(m_table_connections);
+
+    for (auto &table : *m_tables) {
+        if (table.with_id_cache()) {
+            auto &cache = get_id_cache(table);
+            if (get_options()->append) {
+                log_debug("Initializing cache for table '{}' from database...",
+                          table.name());
+                auto const result = m_db_connection.exec(
+                    "SELECT \"{}\" FROM {}", table.id_column_names(),
+                    table.full_name());
+
+                cache.reserve(result.num_tuples());
+                for (int i = 0; i < result.num_tuples(); ++i) {
+                    cache.push_back(
+                        osmium::string_to_object_id(result.get_value(i, 0)));
+                }
+            }
+            cache.sort_unique();
+            log_debug("Cache for table '{}' initialized with {} entries.",
+                      table.name(), cache.size());
+        }
+    }
 }
 
 void output_flex_t::after_ways()
@@ -1205,6 +1285,10 @@ void output_flex_t::relation_modify(osmium::Relation const &rel)
 void output_flex_t::start()
 {
     for (auto &table : m_table_connections) {
+        if (table.table().with_id_cache()) {
+            log_debug("Enable cache for table '{}'.", table.table().name());
+            create_id_cache(table.table());
+        }
         table.start(m_db_connection, get_options()->append);
     }
 
@@ -1218,6 +1302,7 @@ output_flex_t::output_flex_t(output_flex_t const *other,
                              std::shared_ptr<db_copy_thread_t> copy_thread)
 : output_t(other, std::move(mid)), m_locators(other->m_locators),
   m_tables(other->m_tables), m_expire_outputs(other->m_expire_outputs),
+  m_id_caches(other->m_id_caches),
   m_db_connection(get_options()->connection_params, "out.flex.thread"),
   m_stage2_way_ids(other->m_stage2_way_ids),
   m_copy_thread(std::move(copy_thread)), m_lua_state(other->m_lua_state),

src/output-flex.hpp

@@ -165,6 +165,7 @@ class output_flex_t : public output_t
     int app_get_bbox();
 
     int table_insert();
+    int table_in_id_cache();
 
     // Get the flex table that is as first parameter on the Lua stack.
     flex_table_t &get_table_from_param();
@@ -216,6 +217,21 @@ class output_flex_t : public output_t
 
     lua_State *lua_state() noexcept { return m_lua_state.get(); }
 
+    void create_id_cache(flex_table_t const &table)
+    {
+        if (table.num() >= m_id_caches.size()) {
+            m_id_caches.resize(table.num() + 1);
+        }
+        m_id_caches[table.num()] = std::make_shared<idlist_t>();
+    }
+
+    idlist_t &get_id_cache(flex_table_t const &table)
+    {
+        auto& c = m_id_caches[table.num()];
+        assert(c);
+        return *c;
+    }
+
     class way_cache_t
     {
     public:
@@ -273,6 +289,8 @@ class output_flex_t : public output_t
     std::shared_ptr<std::vector<expire_output_t>> m_expire_outputs =
         std::make_shared<std::vector<expire_output_t>>();
 
+    std::vector<std::shared_ptr<idlist_t>> m_id_caches;
+
     std::vector<table_connection_t> m_table_connections;
 
     /// The connection to the database server.
@@ -325,5 +343,6 @@ class output_flex_t : public output_t
 };
 
 int lua_trampoline_table_insert(lua_State *lua_state);
+int lua_trampoline_table_in_id_cache(lua_State *lua_state);
 
 #endif // OSM2PGSQL_OUTPUT_FLEX_HPP