snowflake-ml-python 1.6.0__py3-none-any.whl → 1.6.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (252) hide show
  1. snowflake/cortex/_complete.py +7 -33
  2. snowflake/ml/_internal/env_utils.py +11 -5
  3. snowflake/ml/_internal/exceptions/modeling_error_messages.py +4 -1
  4. snowflake/ml/_internal/telemetry.py +156 -20
  5. snowflake/ml/_internal/utils/identifier.py +48 -11
  6. snowflake/ml/_internal/utils/pkg_version_utils.py +8 -22
  7. snowflake/ml/_internal/utils/snowflake_env.py +23 -13
  8. snowflake/ml/_internal/utils/sql_identifier.py +1 -1
  9. snowflake/ml/_internal/utils/table_manager.py +19 -1
  10. snowflake/ml/_internal/utils/uri.py +2 -2
  11. snowflake/ml/data/_internal/arrow_ingestor.py +66 -10
  12. snowflake/ml/data/data_connector.py +88 -9
  13. snowflake/ml/data/data_ingestor.py +18 -1
  14. snowflake/ml/data/{_internal/ingestor_utils.py → ingestor_utils.py} +5 -1
  15. snowflake/ml/data/torch_utils.py +68 -0
  16. snowflake/ml/dataset/dataset.py +1 -3
  17. snowflake/ml/dataset/dataset_metadata.py +3 -1
  18. snowflake/ml/dataset/dataset_reader.py +9 -3
  19. snowflake/ml/feature_store/examples/airline_features/entities.py +16 -0
  20. snowflake/ml/feature_store/examples/airline_features/features/plane_features.py +31 -0
  21. snowflake/ml/feature_store/examples/airline_features/features/weather_features.py +42 -0
  22. snowflake/ml/feature_store/examples/airline_features/source.yaml +7 -0
  23. snowflake/ml/feature_store/examples/citibike_trip_features/features/station_feature.py +10 -4
  24. snowflake/ml/feature_store/examples/citibike_trip_features/features/trip_feature.py +6 -0
  25. snowflake/ml/feature_store/examples/citibike_trip_features/source.yaml +3 -0
  26. snowflake/ml/feature_store/examples/example_helper.py +69 -31
  27. snowflake/ml/feature_store/examples/new_york_taxi_features/entities.py +3 -3
  28. snowflake/ml/feature_store/examples/new_york_taxi_features/features/{dropoff_features.py → location_features.py} +14 -9
  29. snowflake/ml/feature_store/examples/new_york_taxi_features/features/trip_features.py +36 -0
  30. snowflake/ml/feature_store/examples/new_york_taxi_features/source.yaml +5 -1
  31. snowflake/ml/feature_store/examples/source_data/airline.yaml +4 -0
  32. snowflake/ml/feature_store/examples/source_data/citibike_trips.yaml +1 -1
  33. snowflake/ml/feature_store/examples/wine_quality_features/entities.py +3 -3
  34. snowflake/ml/feature_store/examples/wine_quality_features/features/managed_wine_features.py +13 -6
  35. snowflake/ml/feature_store/examples/wine_quality_features/features/static_wine_features.py +8 -5
  36. snowflake/ml/feature_store/examples/wine_quality_features/source.yaml +3 -0
  37. snowflake/ml/feature_store/feature_store.py +100 -41
  38. snowflake/ml/feature_store/feature_view.py +149 -5
  39. snowflake/ml/fileset/embedded_stage_fs.py +1 -1
  40. snowflake/ml/fileset/fileset.py +1 -1
  41. snowflake/ml/fileset/sfcfs.py +9 -3
  42. snowflake/ml/model/_client/model/model_impl.py +11 -2
  43. snowflake/ml/model/_client/model/model_version_impl.py +186 -20
  44. snowflake/ml/model/_client/ops/model_ops.py +144 -30
  45. snowflake/ml/model/_client/ops/service_ops.py +312 -0
  46. snowflake/ml/model/_client/service/model_deployment_spec.py +94 -0
  47. snowflake/ml/model/_client/service/model_deployment_spec_schema.py +30 -0
  48. snowflake/ml/model/_client/sql/model_version.py +13 -4
  49. snowflake/ml/model/_client/sql/service.py +196 -0
  50. snowflake/ml/model/_deploy_client/image_builds/server_image_builder.py +1 -1
  51. snowflake/ml/model/_deploy_client/snowservice/deploy.py +3 -3
  52. snowflake/ml/model/_model_composer/model_composer.py +5 -0
  53. snowflake/ml/model/_model_composer/model_manifest/model_manifest.py +13 -10
  54. snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py +3 -0
  55. snowflake/ml/model/_packager/model_env/model_env.py +7 -2
  56. snowflake/ml/model/_packager/model_handlers/_base.py +29 -12
  57. snowflake/ml/model/_packager/model_handlers/_utils.py +46 -14
  58. snowflake/ml/model/_packager/model_handlers/catboost.py +25 -16
  59. snowflake/ml/model/_packager/model_handlers/custom.py +6 -2
  60. snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py +32 -20
  61. snowflake/ml/model/_packager/model_handlers/lightgbm.py +23 -56
  62. snowflake/ml/model/_packager/model_handlers/llm.py +11 -5
  63. snowflake/ml/model/_packager/model_handlers/mlflow.py +8 -3
  64. snowflake/ml/model/_packager/model_handlers/model_objective_utils.py +116 -0
  65. snowflake/ml/model/_packager/model_handlers/pytorch.py +8 -3
  66. snowflake/ml/model/_packager/model_handlers/sentence_transformers.py +8 -3
  67. snowflake/ml/model/_packager/model_handlers/sklearn.py +99 -4
  68. snowflake/ml/model/_packager/model_handlers/snowmlmodel.py +123 -5
  69. snowflake/ml/model/_packager/model_handlers/tensorflow.py +9 -4
  70. snowflake/ml/model/_packager/model_handlers/torchscript.py +10 -5
  71. snowflake/ml/model/_packager/model_handlers/xgboost.py +56 -47
  72. snowflake/ml/model/_packager/model_meta/model_meta.py +35 -2
  73. snowflake/ml/model/_packager/model_meta/model_meta_schema.py +11 -0
  74. snowflake/ml/model/_packager/model_packager.py +4 -1
  75. snowflake/ml/model/_packager/model_runtime/model_runtime.py +4 -2
  76. snowflake/ml/model/_signatures/pytorch_handler.py +1 -1
  77. snowflake/ml/model/_signatures/utils.py +9 -0
  78. snowflake/ml/model/models/llm.py +3 -1
  79. snowflake/ml/model/type_hints.py +10 -4
  80. snowflake/ml/modeling/_internal/constants.py +1 -0
  81. snowflake/ml/modeling/_internal/local_implementations/pandas_handlers.py +5 -5
  82. snowflake/ml/modeling/_internal/local_implementations/pandas_trainer.py +9 -6
  83. snowflake/ml/modeling/_internal/model_specifications.py +2 -0
  84. snowflake/ml/modeling/_internal/model_trainer.py +1 -0
  85. snowflake/ml/modeling/_internal/snowpark_implementations/distributed_hpo_trainer.py +2 -2
  86. snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py +5 -5
  87. snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py +113 -160
  88. snowflake/ml/modeling/calibration/calibrated_classifier_cv.py +60 -21
  89. snowflake/ml/modeling/cluster/affinity_propagation.py +60 -21
  90. snowflake/ml/modeling/cluster/agglomerative_clustering.py +60 -21
  91. snowflake/ml/modeling/cluster/birch.py +60 -21
  92. snowflake/ml/modeling/cluster/bisecting_k_means.py +60 -21
  93. snowflake/ml/modeling/cluster/dbscan.py +60 -21
  94. snowflake/ml/modeling/cluster/feature_agglomeration.py +60 -21
  95. snowflake/ml/modeling/cluster/k_means.py +60 -21
  96. snowflake/ml/modeling/cluster/mean_shift.py +60 -21
  97. snowflake/ml/modeling/cluster/mini_batch_k_means.py +60 -21
  98. snowflake/ml/modeling/cluster/optics.py +60 -21
  99. snowflake/ml/modeling/cluster/spectral_biclustering.py +60 -21
  100. snowflake/ml/modeling/cluster/spectral_clustering.py +60 -21
  101. snowflake/ml/modeling/cluster/spectral_coclustering.py +60 -21
  102. snowflake/ml/modeling/compose/column_transformer.py +60 -21
  103. snowflake/ml/modeling/compose/transformed_target_regressor.py +60 -21
  104. snowflake/ml/modeling/covariance/elliptic_envelope.py +60 -21
  105. snowflake/ml/modeling/covariance/empirical_covariance.py +60 -21
  106. snowflake/ml/modeling/covariance/graphical_lasso.py +60 -21
  107. snowflake/ml/modeling/covariance/graphical_lasso_cv.py +60 -21
  108. snowflake/ml/modeling/covariance/ledoit_wolf.py +60 -21
  109. snowflake/ml/modeling/covariance/min_cov_det.py +60 -21
  110. snowflake/ml/modeling/covariance/oas.py +60 -21
  111. snowflake/ml/modeling/covariance/shrunk_covariance.py +60 -21
  112. snowflake/ml/modeling/decomposition/dictionary_learning.py +60 -21
  113. snowflake/ml/modeling/decomposition/factor_analysis.py +60 -21
  114. snowflake/ml/modeling/decomposition/fast_ica.py +60 -21
  115. snowflake/ml/modeling/decomposition/incremental_pca.py +60 -21
  116. snowflake/ml/modeling/decomposition/kernel_pca.py +60 -21
  117. snowflake/ml/modeling/decomposition/mini_batch_dictionary_learning.py +60 -21
  118. snowflake/ml/modeling/decomposition/mini_batch_sparse_pca.py +60 -21
  119. snowflake/ml/modeling/decomposition/pca.py +60 -21
  120. snowflake/ml/modeling/decomposition/sparse_pca.py +60 -21
  121. snowflake/ml/modeling/decomposition/truncated_svd.py +60 -21
  122. snowflake/ml/modeling/discriminant_analysis/linear_discriminant_analysis.py +60 -21
  123. snowflake/ml/modeling/discriminant_analysis/quadratic_discriminant_analysis.py +60 -21
  124. snowflake/ml/modeling/ensemble/ada_boost_classifier.py +60 -21
  125. snowflake/ml/modeling/ensemble/ada_boost_regressor.py +60 -21
  126. snowflake/ml/modeling/ensemble/bagging_classifier.py +60 -21
  127. snowflake/ml/modeling/ensemble/bagging_regressor.py +60 -21
  128. snowflake/ml/modeling/ensemble/extra_trees_classifier.py +60 -21
  129. snowflake/ml/modeling/ensemble/extra_trees_regressor.py +60 -21
  130. snowflake/ml/modeling/ensemble/gradient_boosting_classifier.py +60 -21
  131. snowflake/ml/modeling/ensemble/gradient_boosting_regressor.py +60 -21
  132. snowflake/ml/modeling/ensemble/hist_gradient_boosting_classifier.py +60 -21
  133. snowflake/ml/modeling/ensemble/hist_gradient_boosting_regressor.py +60 -21
  134. snowflake/ml/modeling/ensemble/isolation_forest.py +60 -21
  135. snowflake/ml/modeling/ensemble/random_forest_classifier.py +60 -21
  136. snowflake/ml/modeling/ensemble/random_forest_regressor.py +60 -21
  137. snowflake/ml/modeling/ensemble/stacking_regressor.py +60 -21
  138. snowflake/ml/modeling/ensemble/voting_classifier.py +60 -21
  139. snowflake/ml/modeling/ensemble/voting_regressor.py +60 -21
  140. snowflake/ml/modeling/feature_selection/generic_univariate_select.py +60 -21
  141. snowflake/ml/modeling/feature_selection/select_fdr.py +60 -21
  142. snowflake/ml/modeling/feature_selection/select_fpr.py +60 -21
  143. snowflake/ml/modeling/feature_selection/select_fwe.py +60 -21
  144. snowflake/ml/modeling/feature_selection/select_k_best.py +60 -21
  145. snowflake/ml/modeling/feature_selection/select_percentile.py +60 -21
  146. snowflake/ml/modeling/feature_selection/sequential_feature_selector.py +60 -21
  147. snowflake/ml/modeling/feature_selection/variance_threshold.py +60 -21
  148. snowflake/ml/modeling/framework/base.py +28 -19
  149. snowflake/ml/modeling/gaussian_process/gaussian_process_classifier.py +60 -21
  150. snowflake/ml/modeling/gaussian_process/gaussian_process_regressor.py +60 -21
  151. snowflake/ml/modeling/impute/iterative_imputer.py +60 -21
  152. snowflake/ml/modeling/impute/knn_imputer.py +60 -21
  153. snowflake/ml/modeling/impute/missing_indicator.py +60 -21
  154. snowflake/ml/modeling/kernel_approximation/additive_chi2_sampler.py +60 -21
  155. snowflake/ml/modeling/kernel_approximation/nystroem.py +60 -21
  156. snowflake/ml/modeling/kernel_approximation/polynomial_count_sketch.py +60 -21
  157. snowflake/ml/modeling/kernel_approximation/rbf_sampler.py +60 -21
  158. snowflake/ml/modeling/kernel_approximation/skewed_chi2_sampler.py +60 -21
  159. snowflake/ml/modeling/kernel_ridge/kernel_ridge.py +60 -21
  160. snowflake/ml/modeling/lightgbm/lgbm_classifier.py +60 -21
  161. snowflake/ml/modeling/lightgbm/lgbm_regressor.py +60 -21
  162. snowflake/ml/modeling/linear_model/ard_regression.py +60 -21
  163. snowflake/ml/modeling/linear_model/bayesian_ridge.py +60 -21
  164. snowflake/ml/modeling/linear_model/elastic_net.py +60 -21
  165. snowflake/ml/modeling/linear_model/elastic_net_cv.py +60 -21
  166. snowflake/ml/modeling/linear_model/gamma_regressor.py +60 -21
  167. snowflake/ml/modeling/linear_model/huber_regressor.py +60 -21
  168. snowflake/ml/modeling/linear_model/lars.py +60 -21
  169. snowflake/ml/modeling/linear_model/lars_cv.py +60 -21
  170. snowflake/ml/modeling/linear_model/lasso.py +60 -21
  171. snowflake/ml/modeling/linear_model/lasso_cv.py +60 -21
  172. snowflake/ml/modeling/linear_model/lasso_lars.py +60 -21
  173. snowflake/ml/modeling/linear_model/lasso_lars_cv.py +60 -21
  174. snowflake/ml/modeling/linear_model/lasso_lars_ic.py +60 -21
  175. snowflake/ml/modeling/linear_model/linear_regression.py +60 -21
  176. snowflake/ml/modeling/linear_model/logistic_regression.py +60 -21
  177. snowflake/ml/modeling/linear_model/logistic_regression_cv.py +60 -21
  178. snowflake/ml/modeling/linear_model/multi_task_elastic_net.py +60 -21
  179. snowflake/ml/modeling/linear_model/multi_task_elastic_net_cv.py +60 -21
  180. snowflake/ml/modeling/linear_model/multi_task_lasso.py +60 -21
  181. snowflake/ml/modeling/linear_model/multi_task_lasso_cv.py +60 -21
  182. snowflake/ml/modeling/linear_model/orthogonal_matching_pursuit.py +60 -21
  183. snowflake/ml/modeling/linear_model/passive_aggressive_classifier.py +60 -21
  184. snowflake/ml/modeling/linear_model/passive_aggressive_regressor.py +60 -21
  185. snowflake/ml/modeling/linear_model/perceptron.py +60 -21
  186. snowflake/ml/modeling/linear_model/poisson_regressor.py +60 -21
  187. snowflake/ml/modeling/linear_model/ransac_regressor.py +60 -21
  188. snowflake/ml/modeling/linear_model/ridge.py +60 -21
  189. snowflake/ml/modeling/linear_model/ridge_classifier.py +60 -21
  190. snowflake/ml/modeling/linear_model/ridge_classifier_cv.py +60 -21
  191. snowflake/ml/modeling/linear_model/ridge_cv.py +60 -21
  192. snowflake/ml/modeling/linear_model/sgd_classifier.py +60 -21
  193. snowflake/ml/modeling/linear_model/sgd_one_class_svm.py +60 -21
  194. snowflake/ml/modeling/linear_model/sgd_regressor.py +60 -21
  195. snowflake/ml/modeling/linear_model/theil_sen_regressor.py +60 -21
  196. snowflake/ml/modeling/linear_model/tweedie_regressor.py +60 -21
  197. snowflake/ml/modeling/manifold/isomap.py +60 -21
  198. snowflake/ml/modeling/manifold/mds.py +60 -21
  199. snowflake/ml/modeling/manifold/spectral_embedding.py +60 -21
  200. snowflake/ml/modeling/manifold/tsne.py +60 -21
  201. snowflake/ml/modeling/mixture/bayesian_gaussian_mixture.py +60 -21
  202. snowflake/ml/modeling/mixture/gaussian_mixture.py +60 -21
  203. snowflake/ml/modeling/multiclass/one_vs_one_classifier.py +60 -21
  204. snowflake/ml/modeling/multiclass/one_vs_rest_classifier.py +60 -21
  205. snowflake/ml/modeling/multiclass/output_code_classifier.py +60 -21
  206. snowflake/ml/modeling/naive_bayes/bernoulli_nb.py +60 -21
  207. snowflake/ml/modeling/naive_bayes/categorical_nb.py +60 -21
  208. snowflake/ml/modeling/naive_bayes/complement_nb.py +60 -21
  209. snowflake/ml/modeling/naive_bayes/gaussian_nb.py +60 -21
  210. snowflake/ml/modeling/naive_bayes/multinomial_nb.py +60 -21
  211. snowflake/ml/modeling/neighbors/k_neighbors_classifier.py +60 -21
  212. snowflake/ml/modeling/neighbors/k_neighbors_regressor.py +60 -21
  213. snowflake/ml/modeling/neighbors/kernel_density.py +60 -21
  214. snowflake/ml/modeling/neighbors/local_outlier_factor.py +60 -21
  215. snowflake/ml/modeling/neighbors/nearest_centroid.py +60 -21
  216. snowflake/ml/modeling/neighbors/nearest_neighbors.py +60 -21
  217. snowflake/ml/modeling/neighbors/neighborhood_components_analysis.py +60 -21
  218. snowflake/ml/modeling/neighbors/radius_neighbors_classifier.py +60 -21
  219. snowflake/ml/modeling/neighbors/radius_neighbors_regressor.py +60 -21
  220. snowflake/ml/modeling/neural_network/bernoulli_rbm.py +60 -21
  221. snowflake/ml/modeling/neural_network/mlp_classifier.py +60 -21
  222. snowflake/ml/modeling/neural_network/mlp_regressor.py +60 -21
  223. snowflake/ml/modeling/parameters/disable_model_tracer.py +5 -0
  224. snowflake/ml/modeling/pipeline/pipeline.py +4 -12
  225. snowflake/ml/modeling/preprocessing/polynomial_features.py +60 -21
  226. snowflake/ml/modeling/semi_supervised/label_propagation.py +60 -21
  227. snowflake/ml/modeling/semi_supervised/label_spreading.py +60 -21
  228. snowflake/ml/modeling/svm/linear_svc.py +60 -21
  229. snowflake/ml/modeling/svm/linear_svr.py +60 -21
  230. snowflake/ml/modeling/svm/nu_svc.py +60 -21
  231. snowflake/ml/modeling/svm/nu_svr.py +60 -21
  232. snowflake/ml/modeling/svm/svc.py +60 -21
  233. snowflake/ml/modeling/svm/svr.py +60 -21
  234. snowflake/ml/modeling/tree/decision_tree_classifier.py +60 -21
  235. snowflake/ml/modeling/tree/decision_tree_regressor.py +60 -21
  236. snowflake/ml/modeling/tree/extra_tree_classifier.py +60 -21
  237. snowflake/ml/modeling/tree/extra_tree_regressor.py +60 -21
  238. snowflake/ml/modeling/xgboost/xgb_classifier.py +63 -23
  239. snowflake/ml/modeling/xgboost/xgb_regressor.py +63 -23
  240. snowflake/ml/modeling/xgboost/xgbrf_classifier.py +63 -23
  241. snowflake/ml/modeling/xgboost/xgbrf_regressor.py +63 -23
  242. snowflake/ml/registry/_manager/model_manager.py +20 -2
  243. snowflake/ml/registry/model_registry.py +1 -1
  244. snowflake/ml/registry/registry.py +1 -2
  245. snowflake/ml/utils/sql_client.py +22 -0
  246. snowflake/ml/version.py +1 -1
  247. {snowflake_ml_python-1.6.0.dist-info → snowflake_ml_python-1.6.2.dist-info}/METADATA +55 -3
  248. {snowflake_ml_python-1.6.0.dist-info → snowflake_ml_python-1.6.2.dist-info}/RECORD +251 -238
  249. {snowflake_ml_python-1.6.0.dist-info → snowflake_ml_python-1.6.2.dist-info}/WHEEL +1 -1
  250. snowflake/ml/feature_store/examples/new_york_taxi_features/features/pickup_features.py +0 -58
  251. {snowflake_ml_python-1.6.0.dist-info → snowflake_ml_python-1.6.2.dist-info}/LICENSE.txt +0 -0
  252. {snowflake_ml_python-1.6.0.dist-info → snowflake_ml_python-1.6.2.dist-info}/top_level.txt +0 -0
@@ -52,6 +52,7 @@ from snowflake.ml.feature_store.feature_view import (
52
52
  FeatureViewVersion,
53
53
  _FeatureViewMetadata,
54
54
  )
55
+ from snowflake.ml.utils import sql_client
55
56
  from snowflake.snowpark import DataFrame, Row, Session, functions as F
56
57
  from snowflake.snowpark.exceptions import SnowparkSQLException
57
58
  from snowflake.snowpark.types import (
@@ -94,13 +95,12 @@ class _FeatureStoreObjInfo:
94
95
  return cls(**state_dict) # type: ignore[arg-type]
95
96
 
96
97
 
97
- # TODO: remove "" after dataset is updated
98
98
  class _FeatureStoreObjTypes(Enum):
99
99
  UNKNOWN = "UNKNOWN" # for forward compatibility
100
100
  MANAGED_FEATURE_VIEW = "MANAGED_FEATURE_VIEW"
101
101
  EXTERNAL_FEATURE_VIEW = "EXTERNAL_FEATURE_VIEW"
102
102
  FEATURE_VIEW_REFRESH_TASK = "FEATURE_VIEW_REFRESH_TASK"
103
- TRAINING_DATA = ""
103
+ TRAINING_DATA = "TRAINING_DATA"
104
104
 
105
105
  @classmethod
106
106
  def parse(cls, val: str) -> _FeatureStoreObjTypes:
@@ -140,9 +140,8 @@ _LIST_FEATURE_VIEW_SCHEMA = StructType(
140
140
  )
141
141
 
142
142
 
143
- class CreationMode(Enum):
144
- FAIL_IF_NOT_EXIST = 1
145
- CREATE_IF_NOT_EXIST = 2
143
+ CreationMode = sql_client.CreationOption
144
+ CreationMode.__module__ = __name__
146
145
 
147
146
 
148
147
  @dataclass(frozen=True)
@@ -426,7 +425,9 @@ class FeatureStore:
426
425
 
427
426
  """
428
427
  name = SqlIdentifier(name)
429
- found_rows = self.list_entities().filter(F.col("NAME") == name.resolved()).collect()
428
+ found_rows = (
429
+ self.list_entities().filter(F.col("NAME") == name.resolved()).collect(statement_params=self._telemetry_stmp)
430
+ )
430
431
 
431
432
  if len(found_rows) == 0:
432
433
  warnings.warn(
@@ -603,7 +604,7 @@ class FeatureStore:
603
604
  logger.info(f"Registered FeatureView {feature_view.name}/{version} successfully.")
604
605
  return self.get_feature_view(feature_view.name, str(version))
605
606
 
606
- @dispatch_decorator()
607
+ @overload
607
608
  def update_feature_view(
608
609
  self,
609
610
  name: str,
@@ -612,13 +613,37 @@ class FeatureStore:
612
613
  refresh_freq: Optional[str] = None,
613
614
  warehouse: Optional[str] = None,
614
615
  desc: Optional[str] = None,
616
+ ) -> FeatureView:
617
+ ...
618
+
619
+ @overload
620
+ def update_feature_view(
621
+ self,
622
+ name: FeatureView,
623
+ version: Optional[str] = None,
624
+ *,
625
+ refresh_freq: Optional[str] = None,
626
+ warehouse: Optional[str] = None,
627
+ desc: Optional[str] = None,
628
+ ) -> FeatureView:
629
+ ...
630
+
631
+ @dispatch_decorator() # type: ignore[misc]
632
+ def update_feature_view(
633
+ self,
634
+ name: Union[FeatureView, str],
635
+ version: Optional[str] = None,
636
+ *,
637
+ refresh_freq: Optional[str] = None,
638
+ warehouse: Optional[str] = None,
639
+ desc: Optional[str] = None,
615
640
  ) -> FeatureView:
616
641
  """Update a registered feature view.
617
642
  Check feature_view.py for which fields are allowed to be updated after registration.
618
643
 
619
644
  Args:
620
- name: name of the FeatureView to be updated.
621
- version: version of the FeatureView to be updated.
645
+ name: FeatureView object or name to suspend.
646
+ version: Optional version of feature view. Must set when argument feature_view is a str.
622
647
  refresh_freq: updated refresh frequency.
623
648
  warehouse: updated warehouse.
624
649
  desc: description of feature view.
@@ -660,7 +685,7 @@ class FeatureStore:
660
685
  SnowflakeMLException: [RuntimeError] If FeatureView is not managed and refresh_freq is defined.
661
686
  SnowflakeMLException: [RuntimeError] Failed to update feature view.
662
687
  """
663
- feature_view = self.get_feature_view(name=name, version=version)
688
+ feature_view = self._validate_feature_view_name_and_version_input(name, version)
664
689
  new_desc = desc if desc is not None else feature_view.desc
665
690
 
666
691
  if feature_view.status == FeatureViewStatus.STATIC:
@@ -695,7 +720,7 @@ class FeatureStore:
695
720
  f"Update feature view {feature_view.name}/{feature_view.version} failed: {e}"
696
721
  ),
697
722
  ) from e
698
- return self.get_feature_view(name=name, version=version)
723
+ return self.get_feature_view(name=feature_view.name, version=str(feature_view.version))
699
724
 
700
725
  @overload
701
726
  def read_feature_view(self, feature_view: str, version: str) -> DataFrame:
@@ -853,7 +878,9 @@ class FeatureStore:
853
878
  original_exception=ValueError(f"Failed to find FeatureView {name}/{version}: {results}"),
854
879
  )
855
880
 
856
- return self._compose_feature_view(results[0][0], results[0][1], self.list_entities().collect())
881
+ return self._compose_feature_view(
882
+ results[0][0], results[0][1], self.list_entities().collect(statement_params=self._telemetry_stmp)
883
+ )
857
884
 
858
885
  @overload
859
886
  def refresh_feature_view(self, feature_view: FeatureView) -> None:
@@ -1223,7 +1250,11 @@ class FeatureStore:
1223
1250
  """
1224
1251
  name = SqlIdentifier(name)
1225
1252
  try:
1226
- result = self.list_entities().filter(F.col("NAME") == name.resolved()).collect()
1253
+ result = (
1254
+ self.list_entities()
1255
+ .filter(F.col("NAME") == name.resolved())
1256
+ .collect(statement_params=self._telemetry_stmp)
1257
+ )
1227
1258
  except Exception as e:
1228
1259
  raise snowml_exceptions.SnowflakeMLException(
1229
1260
  error_code=error_codes.INTERNAL_SNOWPARK_ERROR,
@@ -1357,7 +1388,7 @@ class FeatureStore:
1357
1388
  if len(features) == 0:
1358
1389
  raise ValueError("features cannot be empty")
1359
1390
  if isinstance(features[0], str):
1360
- features = self._load_serialized_feature_objects(cast(List[str], features))
1391
+ features = self._load_serialized_feature_views(cast(List[str], features))
1361
1392
 
1362
1393
  df, _ = self._join_features(
1363
1394
  spine_df,
@@ -1441,8 +1472,19 @@ class FeatureStore:
1441
1472
  if save_as is not None:
1442
1473
  try:
1443
1474
  save_as = self._get_fully_qualified_name(save_as)
1444
- result_df.write.mode("errorifexists").save_as_table(save_as)
1475
+ result_df.write.mode("errorifexists").save_as_table(save_as, statement_params=self._telemetry_stmp)
1476
+
1477
+ # Add tag
1478
+ task_obj_info = _FeatureStoreObjInfo(_FeatureStoreObjTypes.TRAINING_DATA, snowml_version.VERSION)
1479
+ self._session.sql(
1480
+ f"""
1481
+ ALTER TABLE {save_as}
1482
+ SET TAG {self._get_fully_qualified_name(_FEATURE_STORE_OBJECT_TAG)}='{task_obj_info.to_json()}'
1483
+ """
1484
+ ).collect(statement_params=self._telemetry_stmp)
1485
+
1445
1486
  return self._session.table(save_as)
1487
+
1446
1488
  except SnowparkSQLException as e:
1447
1489
  if e.sql_error_code == sql_error_codes.OBJECT_ALREADY_EXISTS:
1448
1490
  raise snowml_exceptions.SnowflakeMLException(
@@ -1572,7 +1614,7 @@ class FeatureStore:
1572
1614
 
1573
1615
  fs_meta = FeatureStoreMetadata(
1574
1616
  spine_query=spine_df.queries["queries"][-1],
1575
- serialized_feature_views=[fv.to_json() for fv in features],
1617
+ compact_feature_views=[fv._get_compact_repr().to_json() for fv in features],
1576
1618
  spine_timestamp_col=spine_timestamp_col,
1577
1619
  )
1578
1620
 
@@ -1607,6 +1649,7 @@ class FeatureStore:
1607
1649
  " to generate the data as a Snowflake Table."
1608
1650
  ),
1609
1651
  )
1652
+ # TODO: Add feature store tag once Dataset (version) supports tags
1610
1653
  ds: dataset.Dataset = dataset.create_from_dataframe(
1611
1654
  self._session,
1612
1655
  name,
@@ -1675,11 +1718,18 @@ class FeatureStore:
1675
1718
  if (
1676
1719
  source_meta is None
1677
1720
  or not isinstance(source_meta.properties, FeatureStoreMetadata)
1678
- or source_meta.properties.serialized_feature_views is None
1721
+ or (
1722
+ source_meta.properties.serialized_feature_views is None
1723
+ and source_meta.properties.compact_feature_views is None
1724
+ )
1679
1725
  ):
1680
1726
  raise ValueError(f"Dataset {ds} does not contain valid feature view information.")
1681
1727
 
1682
- return self._load_serialized_feature_objects(source_meta.properties.serialized_feature_views)
1728
+ properties = source_meta.properties
1729
+ if properties.serialized_feature_views:
1730
+ return self._load_serialized_feature_views(properties.serialized_feature_views)
1731
+ else:
1732
+ return self._load_compact_feature_views(properties.compact_feature_views) # type: ignore[arg-type]
1683
1733
 
1684
1734
  @dispatch_decorator()
1685
1735
  def _clear(self, dryrun: bool = True) -> None:
@@ -1700,8 +1750,8 @@ class FeatureStore:
1700
1750
 
1701
1751
  all_fvs_df = self.list_feature_views()
1702
1752
  all_entities_df = self.list_entities()
1703
- all_fvs_rows = all_fvs_df.collect()
1704
- all_entities_rows = all_entities_df.collect()
1753
+ all_fvs_rows = all_fvs_df.collect(statement_params=self._telemetry_stmp)
1754
+ all_entities_rows = all_entities_df.collect(statement_params=self._telemetry_stmp)
1705
1755
 
1706
1756
  if dryrun:
1707
1757
  logger.info(
@@ -1768,6 +1818,7 @@ class FeatureStore:
1768
1818
  {tagging_clause}
1769
1819
  )
1770
1820
  WAREHOUSE = {warehouse}
1821
+ REFRESH_MODE = {feature_view.refresh_mode}
1771
1822
  AS {feature_view.query}
1772
1823
  """
1773
1824
  self._session.sql(query).collect(block=block, statement_params=self._telemetry_stmp)
@@ -1985,7 +2036,7 @@ class FeatureStore:
1985
2036
  MATCH_CONDITION ( spine.ts >= feature.ts )
1986
2037
  ON spine.id = feature.id;
1987
2038
  """
1988
- ).collect()
2039
+ ).collect(statement_params=self._telemetry_stmp)
1989
2040
  except SnowparkSQLException:
1990
2041
  return False
1991
2042
  return result is not None and len(result) == 1
@@ -2094,7 +2145,7 @@ class FeatureStore:
2094
2145
  if "." not in name:
2095
2146
  return f"{self._config.full_schema_path}.{name}"
2096
2147
 
2097
- db_name, schema_name, object_name, _ = identifier.parse_schema_level_object_identifier(name)
2148
+ db_name, schema_name, object_name = identifier.parse_schema_level_object_identifier(name)
2098
2149
  return "{}.{}.{}".format(
2099
2150
  db_name or self._config.database,
2100
2151
  schema_name or self._config.schema,
@@ -2159,11 +2210,7 @@ class FeatureStore:
2159
2210
  if len(fv_maps.keys()) == 0:
2160
2211
  return self._session.create_dataframe([], schema=_LIST_FEATURE_VIEW_SCHEMA)
2161
2212
 
2162
- filters = (
2163
- [lambda d: d["entityName"].startswith(feature_view_name.resolved())] # type: ignore[union-attr]
2164
- if feature_view_name
2165
- else None
2166
- )
2213
+ filters = [lambda d: d["entityName"].startswith(feature_view_name.resolved())] if feature_view_name else None
2167
2214
  res = self._lookup_tagged_objects(self._get_entity_name(entity_name), filters)
2168
2215
 
2169
2216
  output_values: List[List[Any]] = []
@@ -2254,16 +2301,20 @@ class FeatureStore:
2254
2301
  timestamp_col=timestamp_col,
2255
2302
  desc=desc,
2256
2303
  version=version,
2257
- status=FeatureViewStatus(row["scheduling_state"])
2258
- if len(row["scheduling_state"]) > 0
2259
- else FeatureViewStatus.MASKED,
2304
+ status=(
2305
+ FeatureViewStatus(row["scheduling_state"])
2306
+ if len(row["scheduling_state"]) > 0
2307
+ else FeatureViewStatus.MASKED
2308
+ ),
2260
2309
  feature_descs=self._fetch_column_descs("DYNAMIC TABLE", fv_name),
2261
2310
  refresh_freq=row["target_lag"],
2262
2311
  database=self._config.database.identifier(),
2263
2312
  schema=self._config.schema.identifier(),
2264
- warehouse=SqlIdentifier(row["warehouse"], case_sensitive=True).identifier()
2265
- if len(row["warehouse"]) > 0
2266
- else None,
2313
+ warehouse=(
2314
+ SqlIdentifier(row["warehouse"], case_sensitive=True).identifier()
2315
+ if len(row["warehouse"]) > 0
2316
+ else None
2317
+ ),
2267
2318
  refresh_mode=row["refresh_mode"],
2268
2319
  refresh_mode_reason=row["refresh_mode_reason"],
2269
2320
  owner=row["owner"],
@@ -2366,11 +2417,11 @@ class FeatureStore:
2366
2417
  result.append(row)
2367
2418
  return result
2368
2419
 
2369
- def _load_serialized_feature_objects(
2370
- self, serialized_feature_objs: List[str]
2420
+ def _load_serialized_feature_views(
2421
+ self, serialized_feature_views: List[str]
2371
2422
  ) -> List[Union[FeatureView, FeatureViewSlice]]:
2372
2423
  results: List[Union[FeatureView, FeatureViewSlice]] = []
2373
- for obj in serialized_feature_objs:
2424
+ for obj in serialized_feature_views:
2374
2425
  try:
2375
2426
  obj_type = json.loads(obj)[_FEATURE_OBJ_TYPE]
2376
2427
  except Exception as e:
@@ -2384,6 +2435,14 @@ class FeatureStore:
2384
2435
  raise ValueError(f"Unsupported feature object type: {obj_type}")
2385
2436
  return results
2386
2437
 
2438
+ def _load_compact_feature_views(
2439
+ self, compact_feature_views: List[str]
2440
+ ) -> List[Union[FeatureView, FeatureViewSlice]]:
2441
+ results: List[Union[FeatureView, FeatureViewSlice]] = []
2442
+ for obj in compact_feature_views:
2443
+ results.append(FeatureView._load_from_compact_repr(self._session, obj))
2444
+ return results
2445
+
2387
2446
  def _exclude_columns(self, df: DataFrame, exclude_columns: List[str]) -> DataFrame:
2388
2447
  exclude_columns = to_sql_identifiers(exclude_columns) # type: ignore[assignment]
2389
2448
  df_cols = to_sql_identifiers(df.columns)
@@ -2399,12 +2458,12 @@ class FeatureStore:
2399
2458
 
2400
2459
  def _is_dataset_enabled(self) -> bool:
2401
2460
  try:
2402
- self._session.sql(f"SHOW DATASETS IN SCHEMA {self._config.full_schema_path}").collect()
2461
+ self._session.sql(f"SHOW DATASETS IN SCHEMA {self._config.full_schema_path}").collect(
2462
+ statement_params=self._telemetry_stmp
2463
+ )
2403
2464
  return True
2404
- except SnowparkSQLException as e:
2405
- if "'DATASETS' does not exist" in e.message:
2406
- return False
2407
- raise
2465
+ except SnowparkSQLException:
2466
+ return False
2408
2467
 
2409
2468
  def _check_feature_store_object_versions(self) -> None:
2410
2469
  versions = self._collapse_object_versions()
@@ -6,7 +6,7 @@ import warnings
6
6
  from collections import OrderedDict
7
7
  from dataclasses import asdict, dataclass
8
8
  from enum import Enum
9
- from typing import Any, Dict, List, Optional
9
+ from typing import Any, Dict, List, Optional, Union
10
10
 
11
11
  from snowflake.ml._internal.exceptions import (
12
12
  error_codes,
@@ -60,6 +60,29 @@ class _FeatureViewMetadata:
60
60
  return cls(**state_dict)
61
61
 
62
62
 
63
+ @dataclass(frozen=True)
64
+ class _CompactRepresentation:
65
+ """
66
+ A compact representation for FeatureView and FeatureViewSlice, which contains fully qualified name
67
+ and optionally a list of feature indices (None means all features will be included).
68
+ This is to make the metadata much smaller when generating dataset.
69
+ """
70
+
71
+ db: str
72
+ sch: str
73
+ name: str
74
+ version: str
75
+ feature_indices: Optional[List[int]] = None
76
+
77
+ def to_json(self) -> str:
78
+ return json.dumps(asdict(self))
79
+
80
+ @classmethod
81
+ def from_json(cls, json_str: str) -> _CompactRepresentation:
82
+ state_dict = json.loads(json_str)
83
+ return cls(**state_dict)
84
+
85
+
63
86
  class FeatureViewVersion(str):
64
87
  def __new__(cls, version: str) -> FeatureViewVersion:
65
88
  if not _FEATURE_VIEW_VERSION_RE.match(version) or len(version) > _FEATURE_VIEW_VERSION_MAX_LENGTH:
@@ -115,6 +138,19 @@ class FeatureViewSlice:
115
138
  json_dict["feature_view_ref"] = FeatureView.from_json(json_dict["feature_view_ref"], session)
116
139
  return cls(**json_dict)
117
140
 
141
+ def _get_compact_repr(self) -> _CompactRepresentation:
142
+ return _CompactRepresentation(
143
+ db=self.feature_view_ref.database.identifier(), # type: ignore[union-attr]
144
+ sch=self.feature_view_ref.schema.identifier(), # type: ignore[union-attr]
145
+ name=self.feature_view_ref.name.identifier(),
146
+ version=self.feature_view_ref.version, # type: ignore[arg-type]
147
+ feature_indices=self._feature_names_to_indices(),
148
+ )
149
+
150
+ def _feature_names_to_indices(self) -> List[int]:
151
+ name_to_indices_map = {name: idx for idx, name in enumerate(self.feature_view_ref.feature_names)}
152
+ return [name_to_indices_map[n] for n in self.names]
153
+
118
154
 
119
155
  class FeatureView(lineage_node.LineageNode):
120
156
  """
@@ -196,7 +232,7 @@ class FeatureView(lineage_node.LineageNode):
196
232
  self._database: Optional[SqlIdentifier] = None
197
233
  self._schema: Optional[SqlIdentifier] = None
198
234
  self._warehouse: Optional[SqlIdentifier] = SqlIdentifier(warehouse) if warehouse is not None else None
199
- self._refresh_mode: Optional[str] = None
235
+ self._refresh_mode: Optional[str] = _kwargs.get("refresh_mode", "AUTO")
200
236
  self._refresh_mode_reason: Optional[str] = None
201
237
  self._owner: Optional[str] = None
202
238
  self._validate()
@@ -394,6 +430,54 @@ class FeatureView(lineage_node.LineageNode):
394
430
  def feature_descs(self) -> Dict[SqlIdentifier, str]:
395
431
  return self._feature_desc
396
432
 
433
+ def list_columns(self) -> DataFrame:
434
+ """List all columns and their information.
435
+
436
+ Returns:
437
+ A Snowpark DataFrame contains feature information.
438
+
439
+ Example::
440
+
441
+ >>> fs = FeatureStore(...)
442
+ >>> e = Entity("foo", ["id"], desc='my entity')
443
+ >>> fs.register_entity(e)
444
+ <BLANKLINE>
445
+ >>> draft_fv = FeatureView(
446
+ ... name="fv",
447
+ ... entities=[e],
448
+ ... feature_df=self._session.table(<source_table>).select(["NAME", "ID", "TITLE", "AGE", "TS"]),
449
+ ... timestamp_col="ts",
450
+ >>> ).attach_feature_desc({"AGE": "my age", "TITLE": '"my title"'})
451
+ >>> fv = fs.register_feature_view(draft_fv, '1.0')
452
+ <BLANKLINE>
453
+ >>> fv.list_columns().show()
454
+ --------------------------------------------------
455
+ |"NAME" |"CATEGORY" |"DTYPE" |"DESC" |
456
+ --------------------------------------------------
457
+ |NAME |FEATURE |string(64) | |
458
+ |ID |ENTITY |bigint |my entity |
459
+ |TITLE |FEATURE |string(128) |"my title" |
460
+ |AGE |FEATURE |bigint |my age |
461
+ |TS |TIMESTAMP |bigint |NULL |
462
+ --------------------------------------------------
463
+
464
+ """
465
+ session = self._feature_df.session
466
+ rows = []
467
+ for name, type in self._feature_df.dtypes:
468
+ if SqlIdentifier(name) in self.feature_descs:
469
+ desc = self.feature_descs[SqlIdentifier(name)]
470
+ rows.append((name, "FEATURE", type, desc))
471
+ elif SqlIdentifier(name) == self._timestamp_col:
472
+ rows.append((name, "TIMESTAMP", type, None)) # type: ignore[arg-type]
473
+ else:
474
+ for e in self._entities:
475
+ if SqlIdentifier(name) in e.join_keys:
476
+ rows.append((name, "ENTITY", type, e.desc))
477
+ break
478
+
479
+ return session.create_dataframe(rows, schema=["name", "category", "dtype", "desc"])
480
+
397
481
  @property
398
482
  def refresh_freq(self) -> Optional[str]:
399
483
  return self._refresh_freq
@@ -599,12 +683,50 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
599
683
 
600
684
  return fv_dict
601
685
 
602
- def to_df(self, session: Session) -> DataFrame:
686
+ def to_df(self, session: Optional[Session] = None) -> DataFrame:
687
+ """Convert feature view to a Snowpark DataFrame object.
688
+
689
+ Args:
690
+ session: [deprecated] This argument has no effect. No need to pass a session object.
691
+
692
+ Returns:
693
+ A Snowpark Dataframe object contains the information about feature view.
694
+
695
+ Example::
696
+
697
+ >>> fs = FeatureStore(...)
698
+ >>> e = Entity("foo", ["id"], desc='my entity')
699
+ >>> fs.register_entity(e)
700
+ <BLANKLINE>
701
+ >>> draft_fv = FeatureView(
702
+ ... name="fv",
703
+ ... entities=[e],
704
+ ... feature_df=self._session.table(<source_table>).select(["NAME", "ID", "TITLE", "AGE", "TS"]),
705
+ ... timestamp_col="ts",
706
+ >>> ).attach_feature_desc({"AGE": "my age", "TITLE": '"my title"'})
707
+ >>> fv = fs.register_feature_view(draft_fv, '1.0')
708
+ <BLANKLINE>
709
+ >>> fv.to_df().show()
710
+ ----------------------------------------------------------------...
711
+ |"NAME" |"ENTITIES" |"TIMESTAMP_COL" |"DESC" |
712
+ ----------------------------------------------------------------...
713
+ |FV |[ |TS |foobar |
714
+ | | { | | |
715
+ | | "desc": "my entity", | | |
716
+ | | "join_keys": [ | | |
717
+ | | "ID" | | |
718
+ | | ], | | |
719
+ | | "name": "FOO", | | |
720
+ | | "owner": null | | |
721
+ | | } | | |
722
+ | |] | | |
723
+ ----------------------------------------------------------------...
724
+ """
603
725
  values = list(self._to_dict().values())
604
726
  schema = [x.lstrip("_") for x in list(self._to_dict().keys())]
605
727
  values.append(str(FeatureView._get_physical_name(self._name, self._version))) # type: ignore[arg-type]
606
728
  schema.append("physical_name")
607
- return session.create_dataframe([values], schema=schema)
729
+ return self._feature_df.session.create_dataframe([values], schema=schema)
608
730
 
609
731
  def to_json(self) -> str:
610
732
  state_dict = self._to_dict()
@@ -643,6 +765,14 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
643
765
  session=session,
644
766
  )
645
767
 
768
+ def _get_compact_repr(self) -> _CompactRepresentation:
769
+ return _CompactRepresentation(
770
+ db=self.database.identifier(), # type: ignore[union-attr]
771
+ sch=self.schema.identifier(), # type: ignore[union-attr]
772
+ name=self.name.identifier(),
773
+ version=self.version, # type: ignore[arg-type]
774
+ )
775
+
646
776
  @staticmethod
647
777
  def _get_physical_name(fv_name: SqlIdentifier, fv_version: FeatureViewVersion) -> SqlIdentifier:
648
778
  return SqlIdentifier(
@@ -655,9 +785,23 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
655
785
  )
656
786
  )
657
787
 
788
+ @staticmethod
789
+ def _load_from_compact_repr(session: Session, serialized_repr: str) -> Union[FeatureView, FeatureViewSlice]:
790
+ compact_repr = _CompactRepresentation.from_json(serialized_repr)
791
+
792
+ fs = feature_store.FeatureStore(
793
+ session, compact_repr.db, compact_repr.sch, default_warehouse=session.get_current_warehouse()
794
+ )
795
+ fv = fs.get_feature_view(compact_repr.name, compact_repr.version)
796
+
797
+ if compact_repr.feature_indices is not None:
798
+ feature_names = [fv.feature_names[i] for i in compact_repr.feature_indices]
799
+ return fv.slice(feature_names) # type: ignore[no-any-return]
800
+ return fv # type: ignore[no-any-return]
801
+
658
802
  @staticmethod
659
803
  def _load_from_lineage_node(session: Session, name: str, version: str) -> FeatureView:
660
- db_name, feature_store_name, feature_view_name, _ = identifier.parse_schema_level_object_identifier(name)
804
+ db_name, feature_store_name, feature_view_name = identifier.parse_schema_level_object_identifier(name)
661
805
 
662
806
  session_warehouse = session.get_current_warehouse()
663
807
 
@@ -35,7 +35,7 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
35
35
  **kwargs: Any,
36
36
  ) -> None:
37
37
 
38
- (db, schema, object_name, _) = identifier.parse_schema_level_object_identifier(name)
38
+ (db, schema, object_name) = identifier.parse_schema_level_object_identifier(name)
39
39
  self._name = name # TODO: Require or resolve FQN
40
40
  self._domain = domain
41
41
 
@@ -538,7 +538,7 @@ def _validate_target_stage_loc(snowpark_session: snowpark.Session, target_stage_
538
538
  original_exception=fileset_errors.FileSetLocationError('FileSet location should start with "@".'),
539
539
  )
540
540
  try:
541
- db, schema, stage, _ = identifier.parse_schema_level_object_identifier(target_stage_loc[1:])
541
+ db, schema, stage, _ = identifier.parse_snowflake_stage_path(target_stage_loc[1:])
542
542
  if db is None or schema is None:
543
543
  raise ValueError("The stage path should be in the form '@<database>.<schema>.<stage>/*'")
544
544
  df_stages = snowpark_session.sql(f"Show stages like '{stage}' in SCHEMA {db}.{schema}")
@@ -15,6 +15,7 @@ from snowflake.ml._internal.exceptions import (
15
15
  from snowflake.ml._internal.utils import identifier
16
16
  from snowflake.ml.fileset import stage_fs
17
17
  from snowflake.ml.utils import connection_params
18
+ from snowflake.snowpark import context, exceptions as snowpark_exceptions
18
19
 
19
20
  PROTOCOL_NAME = "sfc"
20
21
 
@@ -84,7 +85,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
84
85
  """
85
86
  if kwargs.get(_RECREATE_FROM_SERIALIZED):
86
87
  try:
87
- snowpark_session = self._create_default_session()
88
+ snowpark_session = self._get_default_session()
88
89
  except Exception as e:
89
90
  raise snowml_exceptions.SnowflakeMLException(
90
91
  error_code=error_codes.SNOWML_DESERIALIZATION_FAILED,
@@ -103,7 +104,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
103
104
 
104
105
  super().__init__(**kwargs)
105
106
 
106
- def _create_default_session(self) -> snowpark.Session:
107
+ def _get_default_session(self) -> snowpark.Session:
107
108
  """Create a Snowpark Session from default login options.
108
109
 
109
110
  Returns:
@@ -114,6 +115,11 @@ class SFFileSystem(fsspec.AbstractFileSystem):
114
115
  ValueError: Snowflake Connection could not be created.
115
116
 
116
117
  """
118
+ try:
119
+ return context.get_active_session()
120
+ except snowpark_exceptions.SnowparkSessionException:
121
+ pass
122
+
117
123
  try:
118
124
  snowflake_config = connection_params.SnowflakeLoginOptions()
119
125
  except Exception as e:
@@ -328,7 +334,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
328
334
  ),
329
335
  )
330
336
  try:
331
- res = identifier.parse_schema_level_object_identifier(path[1:])
337
+ res = identifier.parse_snowflake_stage_path(path[1:])
332
338
  if res[1] is None or res[0] is None or (res[3] and not res[3].startswith("/")):
333
339
  raise ValueError("Invalid path. Missing database or schema identifier.")
334
340
  logging.debug(f"Parsed path: {res}")
@@ -5,7 +5,7 @@ import pandas as pd
5
5
  from snowflake.ml._internal import telemetry
6
6
  from snowflake.ml._internal.utils import sql_identifier
7
7
  from snowflake.ml.model._client.model import model_version_impl
8
- from snowflake.ml.model._client.ops import model_ops
8
+ from snowflake.ml.model._client.ops import model_ops, service_ops
9
9
 
10
10
  _TELEMETRY_PROJECT = "MLOps"
11
11
  _TELEMETRY_SUBPROJECT = "ModelManagement"
@@ -19,6 +19,7 @@ class Model:
19
19
  """Model Object containing multiple versions. Mapping to SQL's MODEL object."""
20
20
 
21
21
  _model_ops: model_ops.ModelOperator
22
+ _service_ops: service_ops.ServiceOperator
22
23
  _model_name: sql_identifier.SqlIdentifier
23
24
 
24
25
  def __init__(self) -> None:
@@ -29,17 +30,23 @@ class Model:
29
30
  cls,
30
31
  model_ops: model_ops.ModelOperator,
31
32
  *,
33
+ service_ops: service_ops.ServiceOperator,
32
34
  model_name: sql_identifier.SqlIdentifier,
33
35
  ) -> "Model":
34
36
  self: "Model" = object.__new__(cls)
35
37
  self._model_ops = model_ops
38
+ self._service_ops = service_ops
36
39
  self._model_name = model_name
37
40
  return self
38
41
 
39
42
  def __eq__(self, __value: object) -> bool:
40
43
  if not isinstance(__value, Model):
41
44
  return False
42
- return self._model_ops == __value._model_ops and self._model_name == __value._model_name
45
+ return (
46
+ self._model_ops == __value._model_ops
47
+ and self._service_ops == __value._service_ops
48
+ and self._model_name == __value._model_name
49
+ )
43
50
 
44
51
  @property
45
52
  def name(self) -> str:
@@ -208,6 +215,7 @@ class Model:
208
215
 
209
216
  return model_version_impl.ModelVersion._ref(
210
217
  self._model_ops,
218
+ service_ops=self._service_ops,
211
219
  model_name=self._model_name,
212
220
  version_name=version_id,
213
221
  )
@@ -235,6 +243,7 @@ class Model:
235
243
  return [
236
244
  model_version_impl.ModelVersion._ref(
237
245
  self._model_ops,
246
+ service_ops=self._service_ops,
238
247
  model_name=self._model_name,
239
248
  version_name=version_name,
240
249
  )