Parse decimal with UseArrowNativeType

Author: Todd Meng

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

@@ -32,28 +32,9 @@ internal SparkStatement(SparkConnection connection)
 
         protected override void SetStatementProperties(TExecuteStatementReq statement)
         {
-            // TODO: Ensure this is set dynamically depending on server capabilities.
-            statement.EnforceResultPersistenceMode = false;
-            statement.ResultPersistenceMode = TResultPersistenceMode.ALL_RESULTS;
             // This seems like a good idea to have the server timeout so it doesn't keep processing unnecessarily.
             // Set in combination with a CancellationToken.
             statement.QueryTimeout = QueryTimeoutSeconds;
-            statement.CanReadArrowResult = true;
-
-#pragma warning disable CS0618 // Type or member is obsolete
-            statement.ConfOverlay = SparkConnection.timestampConfig;
-#pragma warning restore CS0618 // Type or member is obsolete
-            statement.UseArrowNativeTypes = new TSparkArrowTypes
-            {
-                TimestampAsArrow = true,
-                DecimalAsArrow = true,
-
-                // set to false so they return as string
-                // otherwise, they return as ARRAY_TYPE but you can't determine
-                // the object type of the items in the array
-                ComplexTypesAsArrow = false,
-                IntervalTypesAsArrow = false,
-            };
         }
 
         public override string AssemblyName => s_assemblyName;

csharp/src/Drivers/Databricks/DatabricksConnection.cs

@@ -55,6 +55,7 @@ internal class DatabricksConnection : SparkHttpConnection
         private bool _useCloudFetch = true;
         private bool _canDecompressLz4 = true;
         private long _maxBytesPerFile = DefaultMaxBytesPerFile;
+        private bool _useArrowNativeTypes = true;
         private const bool DefaultRetryOnUnavailable = true;
         private const int DefaultTemporarilyUnavailableRetryTimeout = 900;
         private bool _useDescTableExtended = true;
@@ -264,6 +265,11 @@ private void ValidateProperties()
         /// </summary>
         internal bool CanDecompressLz4 => _canDecompressLz4;
 
+        /// <summary>
+        /// Gets whether Arrow native types are used.
+        /// </summary>
+        internal bool UseArrowNativeTypes => _useArrowNativeTypes;
+
         /// <summary>
         /// Gets the maximum bytes per file for CloudFetch.
         /// </summary>
@@ -392,10 +398,17 @@ internal override IArrowArrayStream NewReader<T>(T statement, Schema schema, TGe
                 isLz4Compressed = metadataResp.Lz4Compressed;
             }
 
+            // If we have a schema from the metadata response, reparse it with the statement's useArrowNativeTypes flag
+            if (metadataResp?.Schema != null && schema != null)
+            {
+                var schemaParser = new DatabricksSchemaParser(databricksStatement.UseArrowNativeTypes);
+                schema = schemaParser.GetArrowSchema(metadataResp.Schema, DataTypeConversion);
+            }
+
             return new DatabricksCompositeReader(databricksStatement, schema, isLz4Compressed, TlsOptions, _proxyConfigurator);
         }
 
-        internal override SchemaParser SchemaParser => new DatabricksSchemaParser();
+        internal override SchemaParser SchemaParser => new DatabricksSchemaParser(_useArrowNativeTypes);
 
         public override AdbcStatement CreateStatement()
         {
@@ -442,6 +455,7 @@ protected override async Task HandleOpenSessionResponse(TOpenSessionResp? sessio
                 }
                 _enablePKFK = _enablePKFK && FeatureVersionNegotiator.SupportsPKFK(version);
                 _enableMultipleCatalogSupport = session.__isset.canUseMultipleCatalogs ? session.CanUseMultipleCatalogs : false;
+                _useArrowNativeTypes = FeatureVersionNegotiator.SupportsArrowNativeTypes(version);
                 if (session.__isset.initialNamespace)
                 {
                     _defaultNamespace = session.InitialNamespace;

csharp/src/Drivers/Databricks/DatabricksSchemaParser.cs

@@ -25,8 +25,25 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks
 {
     internal class DatabricksSchemaParser : SchemaParser
     {
+        private readonly bool _useArrowNativeTypes;
+
+        public DatabricksSchemaParser(bool useArrowNativeTypes = true)
+        {
+            _useArrowNativeTypes = useArrowNativeTypes;
+        }
+
         public override IArrowType GetArrowType(TPrimitiveTypeEntry thriftType, DataTypeConversion dataTypeConversion)
         {
+            // For decimal types, check if we're using Arrow native types
+            if (thriftType.Type == TTypeId.DECIMAL_TYPE && !_useArrowNativeTypes)
+            {
+                // When decimal is not returned as Arrow, it's returned as a string
+                return StringType.Default;
+            }
+
+            // For future expansion: handle other types that might be affected by useArrowNativeTypes
+            // For example, timestamps, complex types, etc.
+
             return thriftType.Type switch
             {
                 TTypeId.BIGINT_TYPE => Int64Type.Default,

csharp/src/Drivers/Databricks/DatabricksStatement.cs

@@ -68,6 +68,32 @@ protected override void SetStatementProperties(TExecuteStatementReq statement)
         {
             base.SetStatementProperties(statement);
 
+            // Set Databricks-specific statement properties
+            // TODO: Ensure this is set dynamically depending on server capabilities.
+            statement.EnforceResultPersistenceMode = false;
+            statement.ResultPersistenceMode = TResultPersistenceMode.ALL_RESULTS;
+            statement.CanReadArrowResult = true;
+
+#pragma warning disable CS0618 // Type or member is obsolete
+            statement.ConfOverlay = SparkConnection.timestampConfig;
+#pragma warning restore CS0618 // Type or member is obsolete
+
+            // Set UseArrowNativeTypes based on protocol version
+            if (UseArrowNativeTypes)
+            {
+                statement.UseArrowNativeTypes = new TSparkArrowTypes
+                {
+                    TimestampAsArrow = true,
+                    DecimalAsArrow = true,
+
+                    // set to false so they return as string
+                    // otherwise, they return as ARRAY_TYPE but you can't determine
+                    // the object type of the items in the array
+                    ComplexTypesAsArrow = false,
+                    IntervalTypesAsArrow = false,
+                };
+            }
+
             // Set CloudFetch capabilities
             statement.CanDownloadResult = useCloudFetch;
             statement.CanDecompressLZ4Result = canDecompressLz4;
@@ -157,6 +183,12 @@ internal void SetUseCloudFetch(bool useCloudFetch)
         /// </summary>
         public bool CanDecompressLz4 => canDecompressLz4;
 
+        /// <summary>
+        /// Gets whether to use Arrow native types.
+        /// Returns true if the protocol version is greater than or equal to v5.
+        /// </summary>
+        public bool UseArrowNativeTypes => Connection.UseArrowNativeTypes;
+
         /// <summary>
         /// Sets whether the client can decompress LZ4 compressed results.
         /// </summary>

csharp/src/Drivers/Databricks/FeatureVersionNegotiator.cs

@@ -84,6 +84,13 @@ public static bool SupportsPKFK(TProtocolVersion protocolVersion) =>
         * - ArrowMetadata (field is not used)
         */
 
+        /// <summary>
+        /// Gets whether Arrow native types are supported.
+        /// </summary>
+        /// <param name="protocolVersion">The current protocol version.</param>
+        /// <returns>True if Arrow native types are supported; otherwise, false.</returns>
+        public static bool SupportsArrowNativeTypes(TProtocolVersion protocolVersion) =>
+            SupportsProtocolVersion(protocolVersion, TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V5);
 
         #endregion