Skip to content

Commit a28596b

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
fix type codes
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 0dad966 commit a28596b

File tree

3 files changed

+38
-49
lines changed

3 files changed

+38
-49
lines changed

‎src/databricks/sql/backend/sea/backend.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -784,13 +784,14 @@ def get_tables(
784784
assert isinstance(
785785
result, SeaResultSet
786786
), "Expected SeaResultSet from SEA backend"
787-
result.prepare_metadata_columns(MetadataColumnMappings.TABLE_COLUMNS)
788787

789788
# Apply client-side filtering by table_types
790789
from databricks.sql.backend.sea.utils.filters import ResultSetFilter
791790

792791
result = ResultSetFilter.filter_tables_by_type(result, table_types)
793792

793+
result.prepare_metadata_columns(MetadataColumnMappings.TABLE_COLUMNS)
794+
794795
return result
795796

796797
def get_columns(

‎src/databricks/sql/backend/sea/result_set.py‎

Lines changed: 11 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -342,7 +342,7 @@ def _prepare_column_mapping(self) -> None:
342342

343343
def _normalise_arrow_metadata_cols(self, table: "pyarrow.Table") -> "pyarrow.Table":
344344
"""Transform arrow table columns for metadata normalization."""
345-
if not self._metadata_columns:
345+
if not self._metadata_columns or len(table.schema) == 0:
346346
return table
347347

348348
# Reorder columns and add missing ones
@@ -351,50 +351,38 @@ def _normalise_arrow_metadata_cols(self, table: "pyarrow.Table") -> "pyarrow.Tab
351351

352352
for new_idx, result_column in enumerate(self._metadata_columns or []):
353353
old_idx = (
354-
self._column_index_mapping.get(new_idx)
354+
self._column_index_mapping.get(new_idx, None)
355355
if self._column_index_mapping
356356
else None
357357
)
358358

359-
# Get the source data
360-
if old_idx is not None:
361-
column = table.column(old_idx)
362-
values = column.to_pylist()
363-
else:
364-
values = None
365-
366-
# Apply transformation and create column
367-
if values is not None:
368-
column = pyarrow.array(values)
369-
new_columns.append(column)
370-
else:
371-
# Create column with default/transformed values
372-
null_array = pyarrow.nulls(table.num_rows)
373-
new_columns.append(null_array)
359+
column = (
360+
pyarrow.nulls(table.num_rows)
361+
if old_idx is None
362+
else table.column(old_idx)
363+
)
364+
new_columns.append(column)
374365

375366
column_names.append(result_column.thrift_col_name)
376367

377368
return pyarrow.Table.from_arrays(new_columns, names=column_names)
378369

379370
def _normalise_json_metadata_cols(self, rows: List[List[str]]) -> List[List[Any]]:
380371
"""Transform JSON rows for metadata normalization."""
381-
if not self._metadata_columns:
372+
if not self._metadata_columns or len(rows) == 0:
382373
return rows
383374

384375
transformed_rows = []
385376
for row in rows:
386377
new_row = []
387378
for new_idx, result_column in enumerate(self._metadata_columns or []):
388379
old_idx = (
389-
self._column_index_mapping.get(new_idx)
380+
self._column_index_mapping.get(new_idx, None)
390381
if self._column_index_mapping
391382
else None
392383
)
393-
if old_idx is not None:
394-
value = row[old_idx]
395-
else:
396-
value = None
397384

385+
value = None if old_idx is None else row[old_idx]
398386
new_row.append(value)
399387
transformed_rows.append(new_row)
400388
return transformed_rows

‎src/databricks/sql/backend/sea/utils/metadata_mappings.py‎

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -6,32 +6,32 @@ class MetadataColumnMappings:
66
"""Column mappings for metadata queries following JDBC specification."""
77

88
CATALOG_COLUMN_FOR_GET_CATALOGS = ResultColumn(
9-
"TABLE_CAT", "catalog", SqlType.VARCHAR
9+
"TABLE_CAT", "catalog", SqlType.STRING
1010
)
1111

12-
CATALOG_FULL_COLUMN = ResultColumn("TABLE_CATALOG", None, SqlType.VARCHAR)
12+
CATALOG_FULL_COLUMN = ResultColumn("TABLE_CATALOG", None, SqlType.STRING)
1313
SCHEMA_COLUMN_FOR_GET_SCHEMA = ResultColumn(
14-
"TABLE_SCHEM", "databaseName", SqlType.VARCHAR
14+
"TABLE_SCHEM", "databaseName", SqlType.STRING
1515
)
1616

17-
CATALOG_COLUMN = ResultColumn("TABLE_CAT", "catalogName", SqlType.VARCHAR)
18-
SCHEMA_COLUMN = ResultColumn("TABLE_SCHEM", "namespace", SqlType.VARCHAR)
19-
TABLE_NAME_COLUMN = ResultColumn("TABLE_NAME", "tableName", SqlType.VARCHAR)
20-
TABLE_TYPE_COLUMN = ResultColumn("TABLE_TYPE", "tableType", SqlType.VARCHAR)
21-
REMARKS_COLUMN = ResultColumn("REMARKS", "remarks", SqlType.VARCHAR)
22-
TYPE_CATALOG_COLUMN = ResultColumn("TYPE_CAT", None, SqlType.VARCHAR)
23-
TYPE_SCHEM_COLUMN = ResultColumn("TYPE_SCHEM", None, SqlType.VARCHAR)
24-
TYPE_NAME_COLUMN = ResultColumn("TYPE_NAME", None, SqlType.VARCHAR)
17+
CATALOG_COLUMN = ResultColumn("TABLE_CAT", "catalogName", SqlType.STRING)
18+
SCHEMA_COLUMN = ResultColumn("TABLE_SCHEM", "namespace", SqlType.STRING)
19+
TABLE_NAME_COLUMN = ResultColumn("TABLE_NAME", "tableName", SqlType.STRING)
20+
TABLE_TYPE_COLUMN = ResultColumn("TABLE_TYPE", "tableType", SqlType.STRING)
21+
REMARKS_COLUMN = ResultColumn("REMARKS", "remarks", SqlType.STRING)
22+
TYPE_CATALOG_COLUMN = ResultColumn("TYPE_CAT", None, SqlType.STRING)
23+
TYPE_SCHEM_COLUMN = ResultColumn("TYPE_SCHEM", None, SqlType.STRING)
24+
TYPE_NAME_COLUMN = ResultColumn("TYPE_NAME", None, SqlType.STRING)
2525
SELF_REFERENCING_COL_NAME_COLUMN = ResultColumn(
26-
"SELF_REFERENCING_COL_NAME", None, SqlType.VARCHAR
26+
"SELF_REFERENCING_COL_NAME", None, SqlType.STRING
2727
)
28-
REF_GENERATION_COLUMN = ResultColumn("REF_GENERATION", None, SqlType.VARCHAR)
28+
REF_GENERATION_COLUMN = ResultColumn("REF_GENERATION", None, SqlType.STRING)
2929

30-
COL_NAME_COLUMN = ResultColumn("COLUMN_NAME", "col_name", SqlType.VARCHAR)
30+
COL_NAME_COLUMN = ResultColumn("COLUMN_NAME", "col_name", SqlType.STRING)
3131
DATA_TYPE_COLUMN = ResultColumn("DATA_TYPE", None, SqlType.INT)
32-
COLUMN_TYPE_COLUMN = ResultColumn("TYPE_NAME", "columnType", SqlType.VARCHAR)
32+
COLUMN_TYPE_COLUMN = ResultColumn("TYPE_NAME", "columnType", SqlType.STRING)
3333
COLUMN_SIZE_COLUMN = ResultColumn("COLUMN_SIZE", "columnSize", SqlType.INT)
34-
BUFFER_LENGTH_COLUMN = ResultColumn("BUFFER_LENGTH", None, SqlType.INT)
34+
BUFFER_LENGTH_COLUMN = ResultColumn("BUFFER_LENGTH", None, SqlType.TINYINT)
3535

3636
DECIMAL_DIGITS_COLUMN = ResultColumn(
3737
"DECIMAL_DIGITS",
@@ -46,22 +46,22 @@ class MetadataColumnMappings:
4646
)
4747

4848
NULLABLE_COLUMN = ResultColumn("NULLABLE", None, SqlType.INT)
49-
COLUMN_DEF_COLUMN = ResultColumn("COLUMN_DEF", None, SqlType.VARCHAR)
49+
COLUMN_DEF_COLUMN = ResultColumn("COLUMN_DEF", None, SqlType.STRING)
5050
SQL_DATA_TYPE_COLUMN = ResultColumn("SQL_DATA_TYPE", None, SqlType.INT)
5151
SQL_DATETIME_SUB_COLUMN = ResultColumn("SQL_DATETIME_SUB", None, SqlType.INT)
5252
CHAR_OCTET_LENGTH_COLUMN = ResultColumn("CHAR_OCTET_LENGTH", None, SqlType.INT)
53-
IS_NULLABLE_COLUMN = ResultColumn("IS_NULLABLE", "isNullable", SqlType.VARCHAR)
53+
IS_NULLABLE_COLUMN = ResultColumn("IS_NULLABLE", "isNullable", SqlType.STRING)
5454

55-
SCOPE_CATALOG_COLUMN = ResultColumn("SCOPE_CATALOG", None, SqlType.VARCHAR)
56-
SCOPE_SCHEMA_COLUMN = ResultColumn("SCOPE_SCHEMA", None, SqlType.VARCHAR)
57-
SCOPE_TABLE_COLUMN = ResultColumn("SCOPE_TABLE", None, SqlType.VARCHAR)
58-
SOURCE_DATA_TYPE_COLUMN = ResultColumn("SOURCE_DATA_TYPE", None, SqlType.INT)
55+
SCOPE_CATALOG_COLUMN = ResultColumn("SCOPE_CATALOG", None, SqlType.STRING)
56+
SCOPE_SCHEMA_COLUMN = ResultColumn("SCOPE_SCHEMA", None, SqlType.STRING)
57+
SCOPE_TABLE_COLUMN = ResultColumn("SCOPE_TABLE", None, SqlType.STRING)
58+
SOURCE_DATA_TYPE_COLUMN = ResultColumn("SOURCE_DATA_TYPE", None, SqlType.SMALLINT)
5959

6060
IS_AUTO_INCREMENT_COLUMN = ResultColumn(
61-
"IS_AUTOINCREMENT", "isAutoIncrement", SqlType.VARCHAR
61+
"IS_AUTOINCREMENT", "isAutoIncrement", SqlType.STRING
6262
)
6363
IS_GENERATED_COLUMN = ResultColumn(
64-
"IS_GENERATEDCOLUMN", "isGenerated", SqlType.VARCHAR
64+
"IS_GENERATEDCOLUMN", "isGenerated", SqlType.STRING
6565
)
6666

6767
CATALOG_COLUMNS = [CATALOG_COLUMN_FOR_GET_CATALOGS]
@@ -108,5 +108,5 @@ class MetadataColumnMappings:
108108
SCOPE_TABLE_COLUMN,
109109
SOURCE_DATA_TYPE_COLUMN,
110110
IS_AUTO_INCREMENT_COLUMN,
111-
# DO NOT INCLUDE IS_GENERATED_COLUMN - Thrift returns 23 columns
111+
# not including IS_GENERATED_COLUMN of SEA because Thrift does not return an equivalent
112112
]

0 commit comments

Comments
 (0)