|
22 | 22 | from sqlglot.schema import MappingSchema |
23 | 23 | from sqlglot.tokens import Token |
24 | 24 |
|
| 25 | +from sqlmesh import EngineAdapter |
25 | 26 | from sqlmesh.core.constants import MAX_MODEL_DEFINITION_SIZE |
26 | 27 | from sqlmesh.utils.errors import SQLMeshError, ConfigError |
27 | 28 | from sqlmesh.utils.pandas import columns_to_types_from_df |
@@ -1134,31 +1135,45 @@ def select_from_values_for_batch_range( |
1134 | 1135 | batch_start: int, |
1135 | 1136 | batch_end: int, |
1136 | 1137 | alias: str = "t", |
| 1138 | + source_columns: t.Optional[t.Set[str]] = None, |
1137 | 1139 | ) -> exp.Select: |
1138 | | - casted_columns = [ |
1139 | | - exp.alias_(exp.cast(exp.column(column), to=kind), column, copy=False) |
1140 | | - for column, kind in columns_to_types.items() |
1141 | | - ] |
| 1140 | + source_columns = source_columns or set(columns_to_types) |
| 1141 | + source_columns_to_types = EngineAdapter.get_source_columns_to_types( |
| 1142 | + columns_to_types, source_columns |
| 1143 | + ) |
1142 | 1144 |
|
1143 | 1145 | if not values: |
1144 | 1146 | # Ensures we don't generate an empty VALUES clause & forces a zero-row output |
1145 | 1147 | where = exp.false() |
1146 | | - expressions = [tuple(exp.cast(exp.null(), to=kind) for kind in columns_to_types.values())] |
| 1148 | + expressions = [ |
| 1149 | + tuple(exp.cast(exp.null(), to=kind) for kind in source_columns_to_types.values()) |
| 1150 | + ] |
1147 | 1151 | else: |
1148 | 1152 | where = None |
1149 | 1153 | expressions = [ |
1150 | | - tuple(transform_values(v, columns_to_types)) for v in values[batch_start:batch_end] |
| 1154 | + tuple(transform_values(v, source_columns_to_types)) |
| 1155 | + for v in values[batch_start:batch_end] |
1151 | 1156 | ] |
1152 | 1157 |
|
1153 | | - values_exp = exp.values(expressions, alias=alias, columns=columns_to_types) |
| 1158 | + values_exp = exp.values(expressions, alias=alias, columns=source_columns_to_types) |
1154 | 1159 | if values: |
1155 | 1160 | # BigQuery crashes on `SELECT CAST(x AS TIMESTAMP) FROM UNNEST([NULL]) AS x`, but not |
1156 | 1161 | # on `SELECT CAST(x AS TIMESTAMP) FROM UNNEST([CAST(NULL AS TIMESTAMP)]) AS x`. This |
1157 | 1162 | # ensures nulls under the `Values` expression are cast to avoid similar issues. |
1158 | | - for value, kind in zip(values_exp.expressions[0].expressions, columns_to_types.values()): |
| 1163 | + for value, kind in zip( |
| 1164 | + values_exp.expressions[0].expressions, source_columns_to_types.values() |
| 1165 | + ): |
1159 | 1166 | if isinstance(value, exp.Null): |
1160 | 1167 | value.replace(exp.cast(value, to=kind)) |
1161 | 1168 |
|
| 1169 | + casted_columns = [ |
| 1170 | + exp.alias_( |
| 1171 | + exp.cast(exp.column(column) if column in source_columns else exp.Null(), to=kind), |
| 1172 | + column, |
| 1173 | + copy=False, |
| 1174 | + ) |
| 1175 | + for column, kind in columns_to_types.items() |
| 1176 | + ] |
1162 | 1177 | return exp.select(*casted_columns).from_(values_exp, copy=False).where(where, copy=False) |
1163 | 1178 |
|
1164 | 1179 |
|
|
0 commit comments