Skip to content
Open
152 changes: 152 additions & 0 deletions csharp/src/DatabricksStatement.cs
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
using OperationType = AdbcDrivers.Databricks.Telemetry.Proto.Operation.Types.Type;
using Apache.Arrow;
using Apache.Arrow.Adbc;
using Apache.Arrow.Ipc;
using AdbcDrivers.HiveServer2;
using AdbcDrivers.HiveServer2.Hive2;
using AdbcDrivers.HiveServer2.Spark;
Expand Down Expand Up @@ -814,12 +815,163 @@ protected override async Task<QueryResult> GetTablesAsync(CancellationToken canc
// Call the base implementation with the potentially modified catalog name
activity?.AddEvent("statement.get_tables.calling_base_implementation");
QueryResult result = await base.GetTablesAsync(cancellationToken);

// Normalize TABLE_TYPE/REMARKS so the Thrift path matches SEA (issue #527).
// The legacy hive_metastore Thrift response returns REMARKS="UNKNOWN" and an empty
// TABLE_TYPE for tables with no comment. SEA returns REMARKS="" and TABLE_TYPE="TABLE";
// the empty-TABLE_TYPE -> "TABLE" default also matches the official JDBC driver's
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
// MetadataResultSetBuilder (the REMARKS="" parity is with SEA, not a JDBC rule).
result = await NormalizeTablesResultAsync(result, cancellationToken);

activity?.SetTag(SemanticConventions.Db.Response.ReturnedRows, result.RowCount);
activity?.AddEvent("statement.get_tables.complete");
return result;
}, activityName: "GetTables");
}

// Canonical TABLE_TYPE used when the server returns a null/empty classification,
// matching the JDBC driver's MetadataResultSetBuilder default.
private const string DefaultTableType = "TABLE";

// Exact sentinel some legacy Thrift servers (e.g. hive_metastore) emit for REMARKS
// when a table has no comment; SEA returns "" instead. Matched case-sensitively so a
// genuine user comment of "Unknown"/"unknown" is preserved rather than erased.
private const string UnknownRemarksPlaceholder = "UNKNOWN";

/// <summary>
/// Normalizes the TABLE_TYPE and REMARKS columns of a GetTables result so the Thrift
/// path produces the same values as SEA/JDBC (issue #527):
/// - TABLE_TYPE: null/empty defaults to "TABLE".
/// - REMARKS: null/empty or the legacy "UNKNOWN" placeholder defaults to "".
/// Returns the original result unchanged when there is nothing to normalize.
/// </summary>
private static async Task<QueryResult> NormalizeTablesResultAsync(QueryResult result, CancellationToken cancellationToken)
{
if (result.Stream == null)
{
return result;
}

var stream = result.Stream;
Schema schema = stream.Schema;
int tableTypeIndex = schema.GetFieldIndex("TABLE_TYPE");
int remarksIndex = schema.GetFieldIndex("REMARKS");

// Nothing to normalize if neither column is present.
if (tableTypeIndex < 0 && remarksIndex < 0)
{
return result;
}

var batches = new List<RecordBatch>();
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
try
{
while (true)
{
RecordBatch? batch = await stream.ReadNextRecordBatchAsync(cancellationToken);
if (batch == null)
{
break;
}
batches.Add(NormalizeTablesBatch(batch, schema, tableTypeIndex, remarksIndex));
}
}
finally
{
stream.Dispose();
}

int rowCount = 0;
foreach (var batch in batches)
{
rowCount += batch.Length;
}

return new QueryResult(rowCount, new InMemoryArrowStream(schema, batches));
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
}

private static RecordBatch NormalizeTablesBatch(RecordBatch batch, Schema schema, int tableTypeIndex, int remarksIndex)
{
var columns = new IArrowArray[batch.ColumnCount];
for (int i = 0; i < batch.ColumnCount; i++)
{
if (i == tableTypeIndex && batch.Column(i) is StringArray tableTypeArray)
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
{
columns[i] = NormalizeStringColumn(tableTypeArray, DefaultTableType, normalizeUnknown: false);
// The source array is replaced and no longer referenced by the new batch; dispose
// it now rather than leaving its Arrow buffers to the finalizer. Untouched columns
// are reused by reference, so they must NOT be disposed here.
tableTypeArray.Dispose();
}
else if (i == remarksIndex && batch.Column(i) is StringArray remarksArray)
{
columns[i] = NormalizeStringColumn(remarksArray, string.Empty, normalizeUnknown: true);
remarksArray.Dispose();
}
else
{
columns[i] = batch.Column(i);
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
}
}

return new RecordBatch(schema, columns, batch.Length);
}

/// <summary>
/// Builds a normalized copy of a string column. Null/empty values are replaced with
/// <paramref name="defaultValue"/>; when <paramref name="normalizeUnknown"/> is true the
/// exact (case-sensitive) legacy "UNKNOWN" sentinel is also replaced with
/// <paramref name="defaultValue"/>, leaving genuine comments like "Unknown" untouched.
/// </summary>
private static StringArray NormalizeStringColumn(StringArray source, string defaultValue, bool normalizeUnknown)
{
var builder = new StringArray.Builder();
for (int i = 0; i < source.Length; i++)
{
string value = source.IsNull(i) ? string.Empty : source.GetString(i);
if (string.IsNullOrEmpty(value)
|| (normalizeUnknown && string.Equals(value, UnknownRemarksPlaceholder, StringComparison.Ordinal)))
{
value = defaultValue;
}
builder.Append(value);
}
return builder.Build();
}

/// <summary>
/// Minimal in-memory <see cref="IArrowArrayStream"/> that replays a fixed set of record batches.
/// Used to return a normalized GetTables result without re-querying the server.
/// </summary>
private sealed class InMemoryArrowStream : IArrowArrayStream
{
private readonly Schema _schema;
private readonly Queue<RecordBatch> _batches;

public InMemoryArrowStream(Schema schema, IEnumerable<RecordBatch> batches)
{
_schema = schema;
_batches = new Queue<RecordBatch>(batches);
}

public Schema Schema => _schema;

public ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
{
cancellationToken.ThrowIfCancellationRequested();
RecordBatch? batch = _batches.Count > 0 ? _batches.Dequeue() : null;
return new ValueTask<RecordBatch?>(batch);
}

public void Dispose()
{
while (_batches.Count > 0)
{
_batches.Dequeue().Dispose();
}
}
}

/// <summary>
/// Overrides the GetColumnsAsync method to handle the SPARK catalog case.
/// When EnableMultipleCatalogSupport is true:
Expand Down
34 changes: 34 additions & 0 deletions csharp/test/E2E/StatementExecution/SeaMetadataE2ETests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,40 @@ public async Task GetTables_ReturnsAllColumnTypes()
Assert.True(row.ContainsKey("REF_GENERATION"), "Should have REF_GENERATION column");
}

// Regression test for #527 (Thrift vs SEA GetTables value parity).
//
// The legacy hive_metastore catalog's Thrift TGetTablesResp returns
// placeholder values that SEA/JDBC do not:
// - REMARKS = "UNKNOWN" (should normalize to "")
// - TABLE_TYPE = "" (should normalize to "TABLE")
// The Databricks Thrift GetTables path must normalize these so the
// values match SEA. hive_metastore is a built-in catalog on every
// Databricks workspace, so no fixture setup is required.
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
[SkippableFact]
public async Task GetTables_NormalizesRemarksAndTableType_Thrift()
{
SkipIfNotConfigured();
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
using var conn = CreateConnection(new Dictionary<string, string>
{
{ DatabricksParameters.Protocol, "thrift" },
// hive_metastore is a non-default catalog; multi-catalog support
// must be enabled for GetTables to query it.
{ DatabricksParameters.EnableMultipleCatalogSupport, "true" }
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
});
var rows = await ReadMetadata(conn, "GetTables", "hive_metastore", "default");
Assert.True(rows.Count > 0, "hive_metastore.default should expose at least one table");

// REMARKS must default to "" (matching SEA/JDBC), never the legacy
// Thrift placeholder "UNKNOWN".
Assert.DoesNotContain(rows, r => r["REMARKS"] == "UNKNOWN");
Comment thread
peco-review-bot[bot] marked this conversation as resolved.

// TABLE_TYPE must always be a canonical, non-empty classification.
// The server returns "" for hive_metastore managed tables; it must
// be normalized to "TABLE".
Assert.DoesNotContain(rows, r => r["TABLE_TYPE"] == "");
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
Assert.Contains(rows, r => r["TABLE_TYPE"] == "TABLE");
}
Comment thread
peco-review-bot[bot] marked this conversation as resolved.

// --- GetColumnsExtended ---

[SkippableFact]
Expand Down
Loading