Execute the read operation to load data from the offline store.
(self, context: ExecutionContext)
| 59 | self.end_time = end_time |
| 60 | |
| 61 | def execute(self, context: ExecutionContext) -> DAGValue: |
| 62 | """Execute the read operation to load data from the offline store.""" |
| 63 | try: |
| 64 | retrieval_job = create_offline_store_retrieval_job( |
| 65 | data_source=self.source, |
| 66 | column_info=self.column_info, |
| 67 | context=context, |
| 68 | start_time=self.start_time, |
| 69 | end_time=self.end_time, |
| 70 | ) |
| 71 | |
| 72 | if hasattr(retrieval_job, "to_ray_dataset"): |
| 73 | ray_dataset = retrieval_job.to_ray_dataset() |
| 74 | else: |
| 75 | try: |
| 76 | arrow_table = retrieval_job.to_arrow() |
| 77 | ray_wrapper = get_ray_wrapper() |
| 78 | ray_dataset = ray_wrapper.from_arrow(arrow_table) |
| 79 | except Exception: |
| 80 | df = retrieval_job.to_df() |
| 81 | ray_wrapper = get_ray_wrapper() |
| 82 | ray_dataset = ray_wrapper.from_pandas(df) |
| 83 | |
| 84 | field_mapping = getattr(self.source, "field_mapping", None) |
| 85 | if field_mapping: |
| 86 | ray_dataset = apply_field_mapping(ray_dataset, field_mapping) |
| 87 | |
| 88 | return DAGValue( |
| 89 | data=ray_dataset, |
| 90 | format=DAGFormat.RAY, |
| 91 | metadata={ |
| 92 | "source": "offline_store", |
| 93 | "source_type": type(self.source).__name__, |
| 94 | "start_time": self.start_time, |
| 95 | "end_time": self.end_time, |
| 96 | }, |
| 97 | ) |
| 98 | |
| 99 | except Exception as e: |
| 100 | logger.error(f"Ray read node failed: {e}") |
| 101 | raise |
| 102 | |
| 103 | |
| 104 | class RayJoinNode(DAGNode): |