KineticaIngestor.cs

Go to the documentation of this file.

using System;
using System.Collections.Generic;


namespace kinetica
{
    public class KineticaIngestor<T>
    {
        [Serializable]
        public class InsertException<T> : KineticaException
        {
            public Uri url { get; private set; }
            public IList<T> records { get; private set; }
            private string message;

            public InsertException( string msg ) : base( msg ) { }

            internal InsertException( Uri url_, IList<T> records_, string msg ) : base ( msg )
            {
                this.message = msg;
                this.url = url_;
                this.records = records_;
            }

            public override string ToString() { return "InsertException: " + message; }
        }  // end class InsertException




        // KineticaIngestor Members:
        // =========================
        public Kinetica kineticaDB { get; }
        public string table_name { get; }
        public int batch_size { get; }
        public IDictionary<string, string> options { get; }
        //public IReadOnlyDictionary<string, string> options { get; }
        public Int64 count_inserted;
        public Int64 count_updated;
        private KineticaType ktype;
        private Utils.RecordKeyBuilder<T>? primaryKeyBuilder;
        private Utils.RecordKeyBuilder<T>? shardKeyBuilder;
        private IList<int>? routingTable;
        private IList<Utils.WorkerQueue<T>> workerQueues;
        private Random random;


        public KineticaIngestor( Kinetica kdb, string tableName,
                                 int batchSize, KineticaType ktype,
                                 Dictionary<string, string>? options = null,
                                 Utils.WorkerList? workers = null )
        {
            this.kineticaDB = kdb;
            this.table_name = tableName;
            this.ktype = ktype;

            // Validate and save the batch size
            if ( batchSize < 1 )
                throw new KineticaException( $"Batch size must be greater than one; given {batchSize}." );
            this.batch_size = batchSize;

            this.options = options;

            // Set up the primary and shard key builders
            // -----------------------------------------
            this.primaryKeyBuilder = new Utils.RecordKeyBuilder<T>( true,  this.ktype );
            this.shardKeyBuilder   = new Utils.RecordKeyBuilder<T>( false, this.ktype );

            // Based on the Java implementation
            if ( this.primaryKeyBuilder.hasKey() )
            {   // There is a primary key for the given T
                // Now check if there is a distinct shard key
                if ( !this.shardKeyBuilder.hasKey()
                     || this.shardKeyBuilder.hasSameKey( this.primaryKeyBuilder ) )
                    this.shardKeyBuilder = this.primaryKeyBuilder; // no distinct shard key
            }
            else  // there is no primary key for the given T
            {
                this.primaryKeyBuilder = null;

                // Check if there is shard key for T
                if ( !this.shardKeyBuilder.hasKey() )
                    this.shardKeyBuilder = null;
            }  // done setting up the key builders


            // Set up the worker queues
            // -------------------------
            // Do we update records if there are matching primary keys in the
            // database already?
            bool updateOnExistingPk = ( (options != null)
                                           && options.ContainsKey( InsertRecordsRequest<T>.Options.UPDATE_ON_EXISTING_PK )
                                           && options[ InsertRecordsRequest<T>.Options.UPDATE_ON_EXISTING_PK ].Equals( InsertRecordsRequest<T>.Options.TRUE ) );
            // Do T type records have a primary key?
            bool hasPrimaryKey = (this.primaryKeyBuilder != null);
            this.workerQueues = [];
            try
            {
                // If no workers are given, try to get them from Kinetica
                if ( ( workers == null ) || ( workers.Count == 0 ) )
                {
                    workers = new Utils.WorkerList( kdb );
                }

                // If we end up with multiple workers, either given by the
                // user or obtained from Kinetica, then use those
                if ( ( workers != null ) && ( workers.Count > 0 ) )
                {
                    // Add worker queues per worker
                    foreach ( System.Uri workerUrl in workers )
                    {
                        string strWorkerUrl = workerUrl.ToString();
                        strWorkerUrl = strWorkerUrl.EndsWith('/') ? strWorkerUrl[..^1] : strWorkerUrl;
                        string insert_records_worker_url_str = $"{strWorkerUrl}/insert/records";
                        System.Uri url = new( insert_records_worker_url_str );
                        Utils.WorkerQueue<T> worker_queue = new( url, batchSize, hasPrimaryKey, updateOnExistingPk );
                        this.workerQueues.Add( worker_queue );
                    }

                    // Get the worker rank information from Kinetica
                    this.routingTable = kdb.adminShowShards().rank;
                    // Check that enough worker URLs are specified
                    for ( int i = 0; i < routingTable.Count; ++i )
                    {
                        if ( this.routingTable[i] > this.workerQueues.Count )
                            throw new KineticaException( "Not enough worker URLs specified." );
                    }
                }
                else // multihead-ingest is NOT turned on; use the regular Kinetica IP address
                {
                    string strWorkerUrl = kdb.URL.ToString();
                    strWorkerUrl = strWorkerUrl.EndsWith('/') ? strWorkerUrl[..^1] : strWorkerUrl;
                    string insertRecordsUrlStr = $"{strWorkerUrl}/insert/records";
                    System.Uri url = new( insertRecordsUrlStr );
                    Utils.WorkerQueue<T> worker_queue = new( url, batchSize, hasPrimaryKey, updateOnExistingPk );
                    this.workerQueues.Add( worker_queue );
                    this.routingTable = null;
                }
            }
            catch ( Exception ex )
            {
                throw new KineticaException( ex.ToString() );
            }

            // Create the random number generator
            this.random = new( (int) DateTime.Now.Ticks );
        }   // end constructor KineticaIngestor


        public Int64 getCountInserted()
        {
            return System.Threading.Interlocked.Read( ref this.count_inserted );
        }


        public Int64 getCountUpdated()
        {
            return System.Threading.Interlocked.Read( ref this.count_updated );
        }


        public void flush()
        {
            foreach ( Utils.WorkerQueue<T> workerQueue in this.workerQueues )
            {
                // Flush the queue
                IList<T> queue = workerQueue.flush();
                // Actually insert the records
                flush( queue, workerQueue.url );
            }
        }  // end public flush


        private void flush( IList<T> queue, System.Uri url )
        {
            if ( queue.Count == 0 )
                return; // nothing to do since the queue is empty

            try
            {
                // Create the /insert/records request and response objects
                // -------------------------------------------------------
                // Encode the records into binary
                List<byte[]> encodedQueue = [];
                foreach ( var record in queue ) encodedQueue.Add( this.kineticaDB.AvroEncode( record ) );
                RawInsertRecordsRequest request = new( this.table_name, encodedQueue, this.options);

                InsertRecordsResponse response = new();

                if ( url == null )
                {
                    response = this.kineticaDB.insertRecordsRaw( request );
                }
                else
                {
                    response = this.kineticaDB.SubmitRequest<InsertRecordsResponse>( url, request );
                }

                // Save the counts of inserted and updated records
                System.Threading.Interlocked.Add( ref this.count_inserted, response.count_inserted );
                System.Threading.Interlocked.Add( ref this.count_updated, response.count_updated );
            }
            catch ( Exception ex )
            {
                throw new InsertException<T>( url, queue, ex.Message );
            }
        }  // end private flush()



        public void insert( T record )
        {
            // Create the record keys
            Utils.RecordKey? primaryKey = null;  // used to check for uniqueness
            Utils.RecordKey? shardKey = null;    // used to find which worker to send this record to

            // Build the primary key, if any
            if ( this.primaryKeyBuilder != null )
                primaryKey = this.primaryKeyBuilder.build( record );

            // Build the shard/routing key, if any
            if ( this.shardKeyBuilder != null )
                shardKey = this.shardKeyBuilder.build( record );

            // Find out which worker to send the record to; then add the record
            // to the appropriate worker's record queue
            Utils.WorkerQueue<T> workerQueue;
            if ( this.routingTable == null )
            {   // no information regarding multiple workers, so get the first/only one
                workerQueue = this.workerQueues[0];
            }
            else if ( shardKey == null )
            {   // there is no shard/routing key, so get a random worker
                workerQueue = this.workerQueues[ random.Next( this.workerQueues.Count ) ];
            }
            else
            {   // Get the worker based on the sharding/routing key
                int worker_index = shardKey.route( this.routingTable );
                workerQueue = this.workerQueues[worker_index];
            }

            // Insert the record into the queue
            IList<T> queue = workerQueue.insert( record, primaryKey );

            // If inserting the queue resulted in flushing the queue, then flush it
            // properly
            if ( queue != null )
            {
                this.flush( queue, workerQueue.url );
            }
        }  // end insert( record )



        public void insert( IList<T> records)
        {
            // Insert one record at a time
            for ( int i = 0; i < records.Count; ++i )
            {
                try
                {
                    this.insert( records[ i ] );
                }
                catch ( InsertException<T> ex )
                {
                    // Add the remaining records to the insertion exception
                    // record queue
                    IList<T> queue = ex.records;

                    for ( int j = i + 1; j < records.Count; ++j )
                    {
                        queue.Add( records[ j ] );
                    }

                    // Rethrow
                    throw ex;
                }  // end try-catch
            }  // end outer for loop
        }  // end insert( records )



    }  // end class KineticaIngestor<T>




}  // end namespace kinetica