Using a KD Tree in Unity for querying the nearest entity is really slow for some reason

Question

I have an entity manager script that manages and updates all the entities in the world (tick rate is 25 ticks per second). This script uses a KD Tree to get the nearest entity on-demand using a library called KNN, https://github.com/ArthurBrussee/KNN. I based this code on the example code in the library. It works fine with around 3 entities that are querying for the closest entity to them every tick (1/25 of a second) but as soon as I have around 15 or more entities it slows down a LOT (from 50 to 2 FPS). I am using a KD Tree because it is supposed to be really fast for doing this kind of nearest entity calculation but for some reason its really slow here. I'm not really familiar with Jobs and KD trees so I'm not sure what is wrong...

Here is a the code I am using for querying the nearest entity:

public class WorldEntityManager : MonoBehaviour {
    public int EntityTicksPerSecond = 25;
    public int MaxEntities = 10000;

    #region KDArray
    public int QueryK = 5;

    private NativeArray<float3> m_queryPositions;

    private NativeArray<float3> m_points;
    private NativeArray<int> m_results;

    private KnnContainer m_container;

    private NativeArray<RangeQueryResult> m_rangeResults;

    private KnnRebuildJob rebuildJob;
    private JobHandle rebuildHandle;
    #endregion

    [SerializeField]
    private List<Entity> entities;
    private float deltaT;

    public void Init() { // Called once on start, equivalent to the default Start() function
        Debug.Log("Initializing World Entity Manager Subsystem. Target entity ticks per second: " + EntityTicksPerSecond);
        entities = new List<Entity>();
        m_points = new NativeArray<float3>(MaxEntities, Allocator.Persistent);

        // Create a container that accelerates querying for neighbours
        m_container = new KnnContainer(m_points, false, Allocator.Persistent); // Skip building for now. We rebuild every tick

        deltaT = 1f / (float)EntityTicksPerSecond;

        Debug.Log("Successfully initialized World Entity Manager Subsystem");
    }

    public T GetEntityInRange<T>(float3 queryPosition, float radius, Func<T, bool> condition) where T : Entity {
        if (!m_queryPositions.IsCreated || m_queryPositions.Length != 1) {
            if (m_queryPositions.IsCreated) {
                m_queryPositions.Dispose();
                m_results.Dispose();
            }

            m_queryPositions = new NativeArray<float3>(1, Allocator.Persistent);
            m_results = new NativeArray<int>(QueryK, Allocator.Persistent);

            // Initialize all the range query results
            m_rangeResults = new NativeArray<RangeQueryResult>(1, Allocator.Persistent);

            // Each range query result object needs to declare upfront what the maximum number of points in range is
            // Allow for a maximum of 10 results, orig 1024
            m_rangeResults[0] = new RangeQueryResult(5, Allocator.Persistent);

            print("fixing m_queryPositions");
        }

        m_queryPositions[0] = queryPosition;

        // Do a range query
        var query = new QueryRangeBatchJob(m_container, m_queryPositions, radius, m_rangeResults);

        // Schedule query, dependent on the rebuild
        // We're only doing a very limited number of points - so allow each query to have it's own job
        query.ScheduleBatch(1, 1, rebuildHandle).Complete();

        //lockEntityModifications = true;
        var results = m_rangeResults[0];
        for (int i = 0; i < results.Length; i++) {
            try {
                Entity entity = entities[results[i]];
                if (entity == null) {
                    Debug.LogError("Null entity found when range checking. It should've been unregistered but it wasn't!");
                    continue;
                }
                if (entity is T && condition(entity as T)) {
                    return entity as T;
                }
            } catch (ArgumentOutOfRangeException e) {
                Debug.LogWarning("entities: " + entities.Count + " results: " + results.Length);
            }
        }
        //lockEntityModifications = false;

        return null;
    }

    private void rebuildKDTree() {
        // Rebuild our datastructure
        rebuildJob = new KnnRebuildJob(m_container);
        rebuildHandle = rebuildJob.Schedule();
    }

    public void TickWorldEntities() { // Ticks every entity in the world, called by the game manager every tick (1/25 of a second)
        rebuildKDTree();

        for (int i = 0; i < entities.Count; i++) {
            Entity entity = entities[i];
            if (entity == null) {
                Debug.LogError("Null entity found when ticking. It should've been unregistered but it wasn't!");
                continue;
            }
            m_points[i] = entity._transform.position;
            if (entity.ShouldTick) {
                entity.Tick();
            }
        }
    }
}

Hopefully someone can help clarify what is wrong or what I should do to resolve this... Thanks!

Answer 1

If you are getting slowdowns on 15 entities it is not the k-d tree that has a problem, it is something else in your code. Using a k-d tree will change the speed if you're comparing 1000 items, 10000 items or more - the overhead of creating a data structure to sort through 15 elements is never worth it. That said, i understand that you want to at some point scale up.

Also - are you new to unity or are you new to DOTS and trying to mix it up with DOTS? I'm not saying this to make fun of you, i just need to know - to know where to focus.

A couple of things strike me as odd:

• Using a MonoBehavior but not using the Update() or FixedUpdate() hooks.
• Using NativeArray's in a MonoBehavior - and they are all allocated Permanent
• Scheduling Jobs with parameters only to take one batch of one item.

There are some other things, but i'm missing information about the rest of the project so i can't know for sure what is going on.

Others have mentioned it, but take this as good advice - because we have all been there - try run through some tutorials.

Are you trying to use DOTS? John Smith already linked a good tutorial here: https://www.youtube.com/watch?v=IO6_6Y_YUdE

Are you "just" trying to use c# Jobs - consider what should be put in Jobs and what should not - you do not gain much from Jobs if they (some atleast) can't be Scheduled parallel.

This is also a DOTS article, but you can take the Jobs part from it if you don't want to use DOTS: https://github.com/Unity-Technologies/EntityComponentSystemSamples

Good luck!

Answer 2

As someone who was digging into this DOTS KNN library. I have to say, you are not doing it right. You should not use for loop through entity in Monobehavior and tick them. This is just not the practice of DOTS and ECS. Please try to learn ECS+DOTS coding style and use this library in a ISystem/SystemBase, then copy the data to Monobehavior after the job complete if you are doing your stuff in Monobehavior.

A good tutorial for introduction of ECS and DOTS now could be : https://www.youtube.com/watch?v=IO6_6Y_YUdE

Then you can read the example code in the DEMO of this KNN library.

Then to code the KNNClass, you create ISystem for ECS job

public partial struct MyQuerySystem : ISystem{
     public void OnCreate(ref SystemState state) {}
     public void OnDestroy(ref SystemState state) {}
     public void OnUpdate(ref SystemState state) {}
}

for ISystem::OnUpdate you do

     var pointsArray = new NativeArray<float3>(arr.Length, Allocator.Persistent); //points array of all existing points, you fill it with some points coordinates, perhaps copy from Some monobehavior static field, or use a job to copy all entity transform position
    
    var kNNcontainer = new KnnContainer(pointsArray, false, Allocator.TempJob); //a knn container that uses these points, false for not building it immediately, we use a job to build it later

Then, you set the points of all query points:

var queryPoints = new NativeArray<float3>(count, Allocator.TempJob);

And you fill the points somehow with coordinates Then you do build job:

 var rebuild = new KnnRebuildJob(kNNcontainer);
 var jobhandle = rebuild.Schedule();//schedule build job

after building, you query the KNN, which will run after previous build job finish.

 var knnQueryResults = new NativeArray<float3>(count * K, Allocator.TempJob); //K as in K-nearest-neighbour
 var batchSize = 32;
 var jobhandle2 = new QueryKNearestBatchJob(kNNcontainer, queryPoints, knnQueryResults).ScheduleBatch(count, batchSize, jobhandle);
 jobhandle2.complete(); // wait for query finish (very fast for count<10K)

Now you use knnQueryResults for any calculation further, the result is an array size of K * count filled with point index, you access it with query index N and k th neighbour as

var neighbourIndex = knnQueryResults[k+N*K]
var neighbourPoint = pointsArray[neighbourIndex]

ECS is very fast in doing this job, I have AMD5800H and get 15000points * 15000queries at 100fps