csharp/src/Drivers/Databricks/Reader/DatabricksReader.cs

/*
* 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.
*/

using System;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Apache.Arrow.Adbc.Drivers.Apache.Hive2;
using Apache.Arrow.Adbc.Tracing;
using Apache.Arrow.Ipc;
using Apache.Hive.Service.Rpc.Thrift;

namespace Apache.Arrow.Adbc.Drivers.Databricks.Reader
{
    internal sealed class DatabricksReader : BaseDatabricksReader
    {
        List<TSparkArrowBatch>? batches;
        int index;
        IArrowReader? reader;

        public DatabricksReader(IHiveServer2Statement statement, Schema schema, IResponse response, TFetchResultsResp? initialResults, bool isLz4Compressed)
            : base(statement, schema, response, isLz4Compressed)
        {
            // If we have direct results, initialize the batches from them
            if (statement.TryGetDirectResults(this.response, out TSparkDirectResults? directResults))
            {
                this.batches = directResults!.ResultSet.Results.ArrowBatches;
                this.hasNoMoreRows = !directResults.ResultSet.HasMoreRows;
            }
            else if (initialResults != null)
            {
                this.batches = initialResults.Results.ArrowBatches;
                this.hasNoMoreRows = !initialResults.HasMoreRows;
            }
        }

        public override async ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
        {
            return await this.TraceActivityAsync(async activity =>
            {
                ThrowIfDisposed();

                while (true)
                {
                    if (this.reader != null)
                    {
                        RecordBatch? next = await this.reader.ReadNextRecordBatchAsync(cancellationToken);
                        if (next != null)
                        {
                            activity?.AddEvent(SemanticConventions.Messaging.Batch.Response, [new(SemanticConventions.Db.Response.ReturnedRows, next.Length)]);
                            return next;
                        }
                        this.reader = null;
                    }

                    if (this.batches != null && this.index < this.batches.Count)
                    {
                        ProcessFetchedBatches();
                        continue;
                    }

                    this.batches = null;
                    this.index = 0;

                    if (this.hasNoMoreRows)
                    {
                        return null;
                    }
                    // TODO: use an expiring cancellationtoken
                    TFetchResultsReq request = new TFetchResultsReq(this.response.OperationHandle!, TFetchOrientation.FETCH_NEXT, this.statement.BatchSize);

                    // Set MaxBytes from DatabricksStatement
                    if (this.statement is DatabricksStatement databricksStatement)
                    {
                        request.MaxBytes = databricksStatement.MaxBytesPerFetchRequest;
                    }

                    TFetchResultsResp response = await this.statement.Connection.Client!.FetchResults(request, cancellationToken);

                    // Make sure we get the arrowBatches
                    this.batches = response.Results.ArrowBatches;
                    for (int i = 0; i < this.batches.Count; i++)
                    {
                        activity?.AddTag(SemanticConventions.Db.Response.ReturnedRows, this.batches[i].RowCount);
                    }

                    this.hasNoMoreRows = !response.HasMoreRows;
                }
            });
        }

        private void ProcessFetchedBatches()
        {
            var batch = this.batches![this.index];

            // Ensure batch data exists
            if (batch.Batch == null || batch.Batch.Length == 0)
            {
                this.index++;
                return;
            }

            try
            {
                ReadOnlyMemory<byte> dataToUse = new ReadOnlyMemory<byte>(batch.Batch);

                // If LZ4 compression is enabled, decompress the data
                if (isLz4Compressed)
                {
                    dataToUse = Lz4Utilities.DecompressLz4(batch.Batch);
                }

                this.reader = new SingleBatch(ArrowSerializationHelpers.DeserializeRecordBatch(this.schema, dataToUse));
            }
            catch (Exception ex)
            {
                // Create concise error message based on exception type
                string errorMessage = ex switch
                {
                    _ when ex.GetType().Name.Contains("LZ4") => $"Batch {this.index}: LZ4 decompression failed - Data may be corrupted",
                    _ => $"Batch {this.index}: Processing failed - {ex.Message}" // Default case for any other exception
                };
                throw new AdbcException(errorMessage, ex);
            }
            this.index++;
        }

        sealed class SingleBatch : IArrowReader
        {
            private RecordBatch? _recordBatch;

            public SingleBatch(RecordBatch recordBatch) => _recordBatch = recordBatch;

            public ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
            {
                RecordBatch? result = _recordBatch;
                _recordBatch = null;
                return new ValueTask<RecordBatch?>(result);
            }
        }
    }
}