|
23 | 23 | from sqlglot.tokens import Token |
24 | 24 |
|
25 | 25 | from sqlmesh.core.constants import MAX_MODEL_DEFINITION_SIZE |
| 26 | +from sqlmesh.utils import get_source_columns_to_types |
26 | 27 | from sqlmesh.utils.errors import SQLMeshError, ConfigError |
27 | 28 | from sqlmesh.utils.pandas import columns_to_types_from_df |
28 | 29 |
|
@@ -1134,31 +1135,43 @@ 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 = get_source_columns_to_types(columns_to_types, source_columns) |
1142 | 1142 |
|
1143 | 1143 | if not values: |
1144 | 1144 | # Ensures we don't generate an empty VALUES clause & forces a zero-row output |
1145 | 1145 | where = exp.false() |
1146 | | - expressions = [tuple(exp.cast(exp.null(), to=kind) for kind in columns_to_types.values())] |
| 1146 | + expressions = [ |
| 1147 | + tuple(exp.cast(exp.null(), to=kind) for kind in source_columns_to_types.values()) |
| 1148 | + ] |
1147 | 1149 | else: |
1148 | 1150 | where = None |
1149 | 1151 | expressions = [ |
1150 | | - tuple(transform_values(v, columns_to_types)) for v in values[batch_start:batch_end] |
| 1152 | + tuple(transform_values(v, source_columns_to_types)) |
| 1153 | + for v in values[batch_start:batch_end] |
1151 | 1154 | ] |
1152 | 1155 |
|
1153 | | - values_exp = exp.values(expressions, alias=alias, columns=columns_to_types) |
| 1156 | + values_exp = exp.values(expressions, alias=alias, columns=source_columns_to_types) |
1154 | 1157 | if values: |
1155 | 1158 | # BigQuery crashes on `SELECT CAST(x AS TIMESTAMP) FROM UNNEST([NULL]) AS x`, but not |
1156 | 1159 | # on `SELECT CAST(x AS TIMESTAMP) FROM UNNEST([CAST(NULL AS TIMESTAMP)]) AS x`. This |
1157 | 1160 | # 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()): |
| 1161 | + for value, kind in zip( |
| 1162 | + values_exp.expressions[0].expressions, source_columns_to_types.values() |
| 1163 | + ): |
1159 | 1164 | if isinstance(value, exp.Null): |
1160 | 1165 | value.replace(exp.cast(value, to=kind)) |
1161 | 1166 |
|
| 1167 | + casted_columns = [ |
| 1168 | + exp.alias_( |
| 1169 | + exp.cast(exp.column(column) if column in source_columns else exp.Null(), to=kind), |
| 1170 | + column, |
| 1171 | + copy=False, |
| 1172 | + ) |
| 1173 | + for column, kind in columns_to_types.items() |
| 1174 | + ] |
1162 | 1175 | return exp.select(*casted_columns).from_(values_exp, copy=False).where(where, copy=False) |
1163 | 1176 |
|
1164 | 1177 |
|
|
0 commit comments