feat(csharp/src/Drivers/Apache/Spark): low latency test cases

csharp/src/Drivers/Apache/Apache.Arrow.Adbc.Drivers.Apache.csproj

@@ -2,6 +2,7 @@
 
   <PropertyGroup>
     <TargetFrameworks>net472;net6.0</TargetFrameworks>
+    <NoWarn>CS8618</NoWarn>
   </PropertyGroup>
 
   <ItemGroup>

csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs

@@ -16,6 +16,7 @@
 */
 
 using System;
+using System.Collections.Generic;
 using System.Threading.Tasks;
 using Apache.Arrow.Ipc;
 using Apache.Hive.Service.Rpc.Thrift;
@@ -28,6 +29,7 @@ public abstract class HiveServer2Statement : AdbcStatement
         private const int BatchSizeDefault = 50000;
         protected internal HiveServer2Connection connection;
         protected internal TOperationHandle? operationHandle;
+        protected internal TExecuteStatementResp _statementResp;
 
         protected HiveServer2Statement(HiveServer2Connection connection)
         {
@@ -38,20 +40,42 @@ protected virtual void SetStatementProperties(TExecuteStatementReq statement)
         {
         }
 
-        protected abstract IArrowArrayStream NewReader<T>(T statement, Schema schema) where T : HiveServer2Statement;
+        protected abstract IArrowArrayStream NewReader<T>(T statement, Schema schema, TFetchResultsResp? firstResult) where T : HiveServer2Statement;
 
         public override QueryResult ExecuteQuery() => ExecuteQueryAsync().AsTask().Result;
 
         public override UpdateResult ExecuteUpdate() => ExecuteUpdateAsync().Result;
 
         public override async ValueTask<QueryResult> ExecuteQueryAsync()
         {
-            await ExecuteStatementAsync();
-            await PollForResponseAsync();
-            Schema schema = await GetSchemaAsync();
+            var resp =  await ExecuteStatementAsync();
 
-            // TODO: Ensure this is set dynamically based on server capabilities
-            return new QueryResult(-1, NewReader(this, schema));
+            if (resp.OperationHandle.HasResultSet)
+            {
+                if(resp.DirectResults == null)
+                {
+                    throw new InvalidOperationException();
+                }
+
+                if(resp.DirectResults.ResultSetMetadata == null)
+                {
+                    throw new InvalidOperationException();
+                }
+
+                Schema schema = SchemaParser.GetArrowSchema(resp.DirectResults.ResultSetMetadata.Schema);
+                var firstBatch = resp.DirectResults.ResultSet;
+
+                return new QueryResult(-1, NewReader(this, schema, firstBatch));
+
+            }
+            else
+            {
+                await PollForResponseAsync();
+                Schema schema = await GetSchemaAsync();
+
+                // TODO: Ensure this is set dynamically based on server capabilities
+                return new QueryResult(-1, NewReader(this, schema, null));
+            }
         }
 
         public override async Task<UpdateResult> ExecuteUpdateAsync()
@@ -106,10 +130,11 @@ public override void SetOption(string key, string value)
             }
         }
 
-        protected async Task ExecuteStatementAsync()
+        protected async Task<TExecuteStatementResp> ExecuteStatementAsync()
         {
             TExecuteStatementReq executeRequest = new TExecuteStatementReq(this.connection.sessionHandle, this.SqlQuery);
             SetStatementProperties(executeRequest);
+            executeRequest.GetDirectResults = new TSparkGetDirectResults(BatchSizeDefault);
             TExecuteStatementResp executeResponse = await this.connection.Client.ExecuteStatement(executeRequest);
             if (executeResponse.Status.StatusCode == TStatusCode.ERROR_STATUS)
             {
@@ -118,11 +143,14 @@ protected async Task ExecuteStatementAsync()
                     .SetNativeError(executeResponse.Status.ErrorCode);
             }
             this.operationHandle = executeResponse.OperationHandle;
+
+            return executeResponse;
         }
 
         protected async Task PollForResponseAsync()
         {
             TGetOperationStatusResp? statusResponse = null;
+
             do
             {
                 if (statusResponse != null) { await Task.Delay(PollTimeMilliseconds); }

csharp/src/Drivers/Apache/Impala/ImpalaStatement.cs

@@ -16,6 +16,7 @@
 */
 
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -40,7 +41,7 @@ public override object GetValue(IArrowArray arrowArray, int index)
             throw new NotSupportedException();
         }
 
-        protected override IArrowArrayStream NewReader<T>(T statement, Schema schema) => new HiveServer2Reader(statement, schema);
+        protected override IArrowArrayStream NewReader<T>(T statement, Schema schema, TFetchResultsResp? firstBatch) => new HiveServer2Reader(statement, schema);
 
         /// <summary>
         /// Provides the constant string key values to the <see cref="AdbcStatement.SetOption(string, string)" /> method.

csharp/src/Drivers/Apache/Spark/SparkConnection.cs

@@ -70,6 +70,9 @@ public class SparkConnection : HiveServer2Connection
         private readonly Lazy<string> _productVersion;
 
         internal static TSparkGetDirectResults sparkGetDirectResults = new TSparkGetDirectResults(1000);
+        HttpClient _httpClient;
+
+        public String TraceId;
 
         internal static readonly Dictionary<string, string> timestampConfig = new Dictionary<string, string>
         {
@@ -256,6 +259,8 @@ internal SparkConnection(IReadOnlyDictionary<string, string> properties)
             : base(properties)
         {
             _productVersion = new Lazy<string>(() => GetProductVersion(), LazyThreadSafetyMode.PublicationOnly);
+            _httpClient = new HttpClient();
+            TraceId = string.Empty;
         }
 
         protected string ProductVersion => _productVersion.Value;
@@ -278,23 +283,53 @@ protected override async ValueTask<TProtocol> CreateProtocolAsync()
             else
                 token = properties[SparkParameters.Password];
 
-            HttpClient httpClient = new HttpClient();
-            httpClient.BaseAddress = new UriBuilder(Uri.UriSchemeHttps, hostName, -1, path).Uri;
-            httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", token);
-            httpClient.DefaultRequestHeaders.UserAgent.ParseAdd(UserAgent);
-            httpClient.DefaultRequestHeaders.AcceptEncoding.Clear();
-            httpClient.DefaultRequestHeaders.AcceptEncoding.Add(new StringWithQualityHeaderValue("identity"));
-            httpClient.DefaultRequestHeaders.ExpectContinue = false;
-
             TConfiguration config = new TConfiguration();
-
-            ThriftHttpTransport transport = new ThriftHttpTransport(httpClient, config);
+            _httpClient.BaseAddress = new UriBuilder(Uri.UriSchemeHttps, hostName, -1, path).Uri;
+            _httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", token);
+            _httpClient.DefaultRequestHeaders.UserAgent.ParseAdd("SimbaSparkJDBCDriver/2.06.15 Python/PyHive");
+            _httpClient.DefaultRequestHeaders.AcceptEncoding.Clear();
+            _httpClient.DefaultRequestHeaders.AcceptEncoding.Add(new StringWithQualityHeaderValue("identity"));
+            _httpClient.DefaultRequestHeaders.ExpectContinue = false;
+            _httpClient.DefaultRequestHeaders.Add("traceparent", Generate());
+            ThriftHttpTransport transport = new ThriftHttpTransport(_httpClient, config);
             // can switch to the one below if want to use the experimental one with IPeekableTransport
             // ThriftHttpTransport transport = new ThriftHttpTransport(httpClient, config);
             await transport.OpenAsync(CancellationToken.None);
             return new TBinaryProtocol(transport);
         }
 
+        private static readonly char[] Characters = "0123456789abcdef".ToCharArray();
+        private static readonly Random Random = new Random();
+
+        public string Generate()
+        {
+            StringBuilder result = new StringBuilder("00-");
+
+            for (int i = 0; i < 32; i++)
+            {
+                result.Append(Characters[Random.Next(Characters.Length)]);
+            }
+
+            result.Append('-');
+
+            for (int i = 0; i < 16; i++)
+            {
+                result.Append(Characters[Random.Next(Characters.Length)]);
+            }
+
+            result.Append("-01");
+
+            TraceId = result.ToString();
+
+            return TraceId;
+        }
+
+        public void ResetTraceId()
+        {
+            _httpClient.DefaultRequestHeaders.Remove("traceparent");
+            _httpClient.DefaultRequestHeaders.Add("traceparent", Generate());
+        }
+
         protected override TOpenSessionReq CreateSessionRequest()
         {
             return new TOpenSessionReq(TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V7)

csharp/src/Drivers/Apache/Spark/SparkStatement.cs

@@ -32,6 +32,17 @@ internal SparkStatement(SparkConnection connection)
             : base(connection)
         {
         }
+        public override QueryResult ExecuteQuery()
+        {
+            var conn = (connection as SparkConnection);
+
+            if(conn != null)
+            {
+                conn.ResetTraceId();
+            }
+
+            return base.ExecuteQuery();
+        }
 
         protected override void SetStatementProperties(TExecuteStatementReq statement)
         {
@@ -55,7 +66,7 @@ protected override void SetStatementProperties(TExecuteStatementReq statement)
             };
         }
 
-        protected override IArrowArrayStream NewReader<T>(T statement, Schema schema) => new SparkReader(statement, schema);
+        protected override IArrowArrayStream NewReader<T>(T statement, Schema schema, TFetchResultsResp? firstBatch) => new SparkReader(statement, schema, firstBatch);
 
         /// <summary>
         /// Provides the constant string key values to the <see cref="AdbcStatement.SetOption(string, string)" /> method.
@@ -72,11 +83,18 @@ sealed class SparkReader : IArrowArrayStream
             List<TSparkArrowBatch>? batches;
             int index;
             IArrowReader? reader;
+            bool hasMoreRows = false;
 
-            public SparkReader(HiveServer2Statement statement, Schema schema)
+            public SparkReader(HiveServer2Statement statement, Schema schema, TFetchResultsResp? firstBatch)
             {
                 this.statement = statement;
                 this.schema = schema;
+
+                if (firstBatch != null)
+                {
+                    this.batches = firstBatch.Results.ArrowBatches;
+                    this.hasMoreRows = firstBatch.HasMoreRows;
+                }
             }
 
             public Schema Schema { get { return schema; } }
@@ -104,7 +122,7 @@ public SparkReader(HiveServer2Statement statement, Schema schema)
                     this.batches = null;
                     this.index = 0;
 
-                    if (this.statement == null)
+                    if (this.statement == null || !hasMoreRows)
                     {
                         return null;
                     }
@@ -117,6 +135,8 @@ public SparkReader(HiveServer2Statement statement, Schema schema)
                     {
                         this.statement = null;
                     }
+
+                    hasMoreRows = response.HasMoreRows;
                 }
             }
 

csharp/test/Drivers/Apache/Spark/DriverTests.cs

@@ -16,10 +16,15 @@
 */
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
+using System.IO;
 using System.Linq;
 using System.Text.RegularExpressions;
+using System.Threading;
 using System.Threading.Tasks;
+using Apache.Arrow.Adbc.Drivers.Apache.Spark;
 using Apache.Arrow.Adbc.Tests.Metadata;
 using Apache.Arrow.Adbc.Tests.Xunit;
 using Apache.Arrow.Ipc;
@@ -505,14 +510,71 @@ public async Task CanGetTableTypes()
         [SkippableFact, Order(10)]
         public void CanExecuteQuery()
         {
-            using AdbcConnection adbcConnection = NewConnection();
+            long count = 0;
+            long target = 10;
+            Queue<String> log = new Queue<string>();
+            string FileName = $@"c:\temp\hmsmeta\{DateTime.Now.ToString("yyyyMMdd_HHmm")}__{target}.log";
 
-            using AdbcStatement statement = adbcConnection.CreateStatement();
-            statement.SqlQuery = TestConfiguration.Query;
-
-            QueryResult queryResult = statement.ExecuteQuery();
+            for (int i = 0; i < target; i++)
+            {
+                ThreadPool.QueueUserWorkItem((_) =>
+                {
+                    using AdbcConnection adbcConnection = NewConnection();
+                    var sparkConn = adbcConnection as SparkConnection;
+                    if(sparkConn == null)
+                    {
+                        throw new InvalidCastException();
+                    }
+
+                    for (int i = 0; i < 10000; i++)
+                    {
+                        var sw = new Stopwatch();
+                        sw.Start();
+                        bool succeed = false;
+                        try
+                        {
+                            using AdbcStatement statement = adbcConnection.CreateStatement();
+                            statement.SqlQuery = TestConfiguration.Query;
+
+                            QueryResult queryResult = statement.ExecuteQuery();
+                            Tests.DriverTests.CanExecuteQuery(queryResult, TestConfiguration.ExpectedResultsCount);
+                            succeed = true;
+                        }
+                        catch (Exception)
+                        {
+                            succeed = false;
+                        }
+                        sw.Stop();
+                        var logStr = $"{DateTime.Now.ToString("yyyy-MM-dd HH:mm:ss.fff")},{sparkConn.TraceId},{sw.ElapsedMilliseconds},{succeed}";
+
+                        lock(log)
+                        {
+                            log.Enqueue(logStr);
+                            if(log.Count > 1000)
+                            {
+                                File.AppendAllLines(FileName, log.ToArray());
+                                log.Clear();
+                            }
+                        }
+                    }
+
+                    lock (log)
+                    {
+                        if (log.Count > 0)
+                        {
+                            File.AppendAllLines(FileName, log.ToArray());
+                            log.Clear();
+                        }
+                    }
+
+                    Interlocked.Increment(ref count);
+                });
+            }
 
-            Tests.DriverTests.CanExecuteQuery(queryResult, TestConfiguration.ExpectedResultsCount);
+            while(Interlocked.Read(ref count) < target)
+            {
+                Thread.Sleep(1000);
+            }
         }
 
         /// <summary>