From 4d58faadfb1382d4cee4808d0a5d9b168321ffd8 Mon Sep 17 00:00:00 2001 From: Jayant Singh Date: Wed, 9 Apr 2025 18:00:59 +0530 Subject: [PATCH 1/2] Improvements to fetch results #1 Signed-off-by: Jayant Singh --- src/databricks/sql/utils.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 348c268da..0fd2cbdbf 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -611,21 +611,32 @@ def convert_arrow_based_set_to_arrow_table(arrow_batches, lz4_compressed, schema def convert_decimals_in_arrow_table(table, description) -> "pyarrow.Table": + new_columns = [] + new_fields = [] + for i, col in enumerate(table.itercolumns()): + field = table.field(i) + if description[i][1] == "decimal": - decimal_col = col.to_pandas().apply( - lambda v: v if v is None else Decimal(v) - ) precision, scale = description[i][4], description[i][5] assert scale is not None assert precision is not None - # Spark limits decimal to a maximum scale of 38, - # so 128 is guaranteed to be big enough + # create the target decimal type dtype = pyarrow.decimal128(precision, scale) - col_data = pyarrow.array(decimal_col, type=dtype) - field = table.field(i).with_type(dtype) - table = table.set_column(i, field, col_data) - return table + + # convert the column directly using PyArrow's cast operation + new_col = col.cast(dtype) + new_field = field.with_type(dtype) + + new_columns.append(new_col) + new_fields.append(new_field) + else: + new_columns.append(col) + new_fields.append(field) + + new_schema = pyarrow.schema(new_fields) + + return pyarrow.Table.from_arrays(new_columns, schema=new_schema) def convert_to_assigned_datatypes_in_column_table(column_table, description): From 5ea3df4d65836a0ddc741528df76f767647772ff Mon Sep 17 00:00:00 2001 From: Jayant Singh Date: Mon, 12 May 2025 08:49:59 +0530 Subject: [PATCH 2/2] nit Signed-off-by: Jayant Singh --- src/databricks/sql/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 0fd2cbdbf..186f13dd6 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -624,7 +624,6 @@ def convert_decimals_in_arrow_table(table, description) -> "pyarrow.Table": # create the target decimal type dtype = pyarrow.decimal128(precision, scale) - # convert the column directly using PyArrow's cast operation new_col = col.cast(dtype) new_field = field.with_type(dtype)