diff --git a/.prow/config.yaml b/.prow/config.yaml index 41f95180fbb..4b6e352a12f 100644 --- a/.prow/config.yaml +++ b/.prow/config.yaml @@ -145,18 +145,18 @@ presubmits: postsubmits: gojek/feast: - name: publish-python-sdk - decorate: true + decorate: true spec: containers: - image: python:3 command: - sh - - -c + - -c - | .prow/scripts/publish-python-sdk.sh \ --directory-path sdk/python --repository pypi volumeMounts: - - name: pypirc + - name: pypirc mountPath: /root/.pypirc subPath: .pypirc readOnly: true @@ -170,7 +170,7 @@ postsubmits: - ^v(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(-(0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(\.(0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*)?(\+[0-9a-zA-Z-]+(\.[0-9a-zA-Z-]+)*)?$ - name: publish-docker-images - decorate: true + decorate: true spec: containers: - image: google/cloud-sdk:273.0.0 @@ -182,14 +182,14 @@ postsubmits: --archive-uri gs://feast-templocation-kf-feast/.m2.2019-10-24.tar \ --output-dir $PWD/ - if [ $PULL_BASE_REF == "master" ]; then - + if [ $PULL_BASE_REF == "master" ]; then + .prow/scripts/publish-docker-image.sh \ --repository gcr.io/kf-feast/feast-core \ --tag dev \ --file infra/docker/core/Dockerfile \ --google-service-account-file /etc/gcloud/service-account.json - + .prow/scripts/publish-docker-image.sh \ --repository gcr.io/kf-feast/feast-serving \ --tag dev \ @@ -203,13 +203,13 @@ postsubmits: docker push gcr.io/kf-feast/feast-serving:${PULL_BASE_SHA} else - + .prow/scripts/publish-docker-image.sh \ --repository gcr.io/kf-feast/feast-core \ --tag ${PULL_BASE_REF:1} \ --file infra/docker/core/Dockerfile \ --google-service-account-file /etc/gcloud/service-account.json - + .prow/scripts/publish-docker-image.sh \ --repository gcr.io/kf-feast/feast-serving \ --tag ${PULL_BASE_REF:1} \ @@ -244,7 +244,7 @@ postsubmits: - ^v(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(-(0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*)(\.(0|[1-9]\d*|\d*[a-zA-Z-][0-9a-zA-Z-]*))*)?(\+[0-9a-zA-Z-]+(\.[0-9a-zA-Z-]+)*)?$ - name: publish-helm-chart - decorate: true + decorate: true spec: containers: - image: google/cloud-sdk:273.0.0-slim @@ -253,7 +253,7 @@ postsubmits: - -c - | gcloud auth activate-service-account --key-file /etc/gcloud/service-account.json - + curl -s https://get.helm.sh/helm-v2.16.1-linux-amd64.tar.gz | tar -C /tmp -xz mv /tmp/linux-amd64/helm /usr/bin/helm helm init --client-only diff --git a/sdk/__init__.py b/sdk/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/sdk/python/feast/client.py b/sdk/python/feast/client.py index 3af2e12a91e..0254f4187ba 100644 --- a/sdk/python/feast/client.py +++ b/sdk/python/feast/client.py @@ -12,13 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. - +import json import logging import os -import sys +import time from collections import OrderedDict +from math import ceil from typing import Dict, Union from typing import List +from urllib.parse import urlparse + +import fastavro import grpc import pandas as pd import pyarrow as pa @@ -33,11 +37,12 @@ GetFeatureSetResponse, ) from feast.core.CoreService_pb2_grpc import CoreServiceStub -from feast.exceptions import format_grpc_exception from feast.feature_set import FeatureSet, Entity from feast.job import Job -from feast.loaders.file import export_dataframe_to_staging_location -from feast.loaders.ingest import ingest_table_to_kafka +from feast.loaders.abstract_producer import get_producer +from feast.loaders.file import export_source_to_staging_location +from feast.loaders.ingest import KAFKA_CHUNK_PRODUCTION_TIMEOUT +from feast.loaders.ingest import get_feature_row_chunks from feast.serving.ServingService_pb2 import GetFeastServingInfoResponse from feast.serving.ServingService_pb2 import ( GetOnlineFeaturesRequest, @@ -257,7 +262,7 @@ def _apply_feature_set(self, feature_set: FeatureSet): print(f"No change detected or applied: {feature_set.name}") # Deep copy from the returned feature set to the local feature set - feature_set.update_from_feature_set(applied_fs) + feature_set._update_from_feature_set(applied_fs) def list_feature_sets(self) -> List[FeatureSet]: """ @@ -319,22 +324,32 @@ def list_entities(self) -> Dict[str, Entity]: return entities_dict def get_batch_features( - self, feature_ids: List[str], entity_rows: pd.DataFrame + self, feature_ids: List[str], entity_rows: Union[pd.DataFrame, str] ) -> Job: """ Retrieves historical features from a Feast Serving deployment. Args: - feature_ids: List of feature ids that will be returned for each - entity. Each feature id should have the following format + feature_ids (List[str]): + List of feature ids that will be returned for each entity. + Each feature id should have the following format "feature_set_name:version:feature_name". - entity_rows: Pandas dataframe containing entities and a 'datetime' - column. Each entity in a feature set must be present as a column - in this dataframe. The datetime column must + + entity_rows (Union[pd.DataFrame, str]): + Either: + Pandas dataframe containing entities and a 'datetime' column. + Each entity in a feature set must be present as a column in this + dataframe. The datetime column must contain timestamps in + datetime64 format. + + Or: + A file path in AVRO format representing the entity rows. Returns: - Returns a job object that can be used to monitor retrieval progress - asynchronously, and can be used to materialize the results + feast.job.Job: + Returns a job object that can be used to monitor retrieval + progress asynchronously, and can be used to materialize the + results. Examples: >>> from feast import Client @@ -357,21 +372,11 @@ def get_batch_features( fs_request = _build_feature_set_request(feature_ids) - # Validate entity rows based on entities in Feast Core - self._validate_entity_rows_for_batch_retrieval(entity_rows, fs_request) - - # Remove timezone from datetime column - if isinstance( - entity_rows["datetime"].dtype, pd.core.dtypes.dtypes.DatetimeTZDtype - ): - entity_rows["datetime"] = pd.DatetimeIndex( - entity_rows["datetime"] - ).tz_localize(None) - # Retrieve serving information to determine store type and # staging location serving_info = self._serving_service_stub.GetFeastServingInfo( - GetFeastServingInfoRequest(), timeout=GRPC_CONNECTION_TIMEOUT_DEFAULT + GetFeastServingInfoRequest(), + timeout=GRPC_CONNECTION_TIMEOUT_DEFAULT ) # type: GetFeastServingInfoResponse if serving_info.type != FeastServingType.FEAST_SERVING_TYPE_BATCH: @@ -380,17 +385,50 @@ def get_batch_features( f"does not support batch retrieval " ) - # Export and upload entity row dataframe to staging location + if isinstance(entity_rows, pd.DataFrame): + # Pandas DataFrame detected + # Validate entity rows to based on entities in Feast Core + self._validate_dataframe_for_batch_retrieval( + entity_rows=entity_rows, + feature_sets_request=fs_request + ) + + # Remove timezone from datetime column + if isinstance( + entity_rows["datetime"].dtype, + pd.core.dtypes.dtypes.DatetimeTZDtype + ): + entity_rows["datetime"] = pd.DatetimeIndex( + entity_rows["datetime"] + ).tz_localize(None) + elif isinstance(entity_rows, str): + # String based source + if entity_rows.endswith((".avro", "*")): + # Validate Avro entity rows to based on entities in Feast Core + self._validate_avro_for_batch_retrieval( + source=entity_rows, + feature_sets_request=fs_request + ) + else: + raise Exception( + f"Only .avro and wildcard paths are accepted as entity_rows" + ) + else: + raise Exception(f"Only pandas.DataFrame and str types are allowed" + f" as entity_rows, but got {type(entity_rows)}.") + + # Export and upload entity row DataFrame to staging location # provided by Feast - staged_file = export_dataframe_to_staging_location( + staged_files = export_source_to_staging_location( entity_rows, serving_info.job_staging_location - ) # type: str + ) # type: List[str] request = GetBatchFeaturesRequest( feature_sets=fs_request, dataset_source=DatasetSource( file_source=DatasetSource.FileSource( - file_uris=[staged_file], data_format=DataFormat.DATA_FORMAT_AVRO + file_uris=staged_files, + data_format=DataFormat.DATA_FORMAT_AVRO ) ), ) @@ -399,28 +437,107 @@ def get_batch_features( response = self._serving_service_stub.GetBatchFeatures(request) return Job(response.job, self._serving_service_stub) - def _validate_entity_rows_for_batch_retrieval( - self, entity_rows, feature_sets_request + def _validate_dataframe_for_batch_retrieval( + self, entity_rows: pd.DataFrame, feature_sets_request ): """ - Validate whether an entity_row dataframe contains the correct - information for batch retrieval + Validate whether an the entity rows in a DataFrame contains the correct + information for batch retrieval. + + Datetime column must be present in the DataFrame. Args: - entity_rows: Pandas dataframe containing entities and datetime - column. Each entity in a feature set must be present as a - column in this dataframe. - feature_sets_request: Feature sets that will be requested + entity_rows (pd.DataFrame): + Pandas DataFrame containing entities and datetime column. Each + entity in a feature set must be present as a column in this + DataFrame. + + feature_sets_request: + Feature sets that will be requested. """ + self._validate_columns( + columns=entity_rows.columns, + feature_sets_request=feature_sets_request, + datetime_field="datetime" + ) + + def _validate_avro_for_batch_retrieval( + self, source: str, feature_sets_request + ): + """ + Validate whether the entity rows in an Avro source file contains the + correct information for batch retrieval. + + Only gs:// and local files (file://) uri schemes are allowed. + + Avro file must have a column named "event_timestamp". + + No checks will be done if a GCS path is provided. + + Args: + source (str): + File path to Avro. + + feature_sets_request: + Feature sets that will be requested. + """ + p = urlparse(source) + + if p.scheme == "gs": + # GCS path provided (Risk is delegated to user) + # No validation if GCS path is provided + return + elif p.scheme == "file" or not p.scheme: + # Local file (file://) provided + file_path = os.path.abspath(os.path.join(p.netloc, p.path)) + else: + raise Exception(f"Unsupported uri scheme provided {p.scheme}, only " + f"local files (file://), and gs:// schemes are " + f"allowed") + + with open(file_path, "rb") as f: + reader = fastavro.reader(f) + schema = json.loads(reader.metadata["avro.schema"]) + columns = [x["name"] for x in schema["fields"]] + self._validate_columns( + columns=columns, + feature_sets_request=feature_sets_request, + datetime_field="event_timestamp" + ) + + def _validate_columns( + self, columns: List[str], + feature_sets_request, + datetime_field: str + ) -> None: + """ + Check if the required column contains the correct values for batch + retrieval. + + Args: + columns (List[str]): + List of columns to validate against feature_sets_request. + + feature_sets_request (): + Feature sets that will be requested. + + datetime_field (str): + Name of the datetime field that must be enforced and present as + a column in the data source. + + Returns: + None: + None + """ # Ensure datetime column exists - if "datetime" not in entity_rows.columns: + if datetime_field not in columns: raise ValueError( - f'Entity rows does not contain "datetime" column in columns ' - f"{entity_rows.columns}" + f'Entity rows does not contain "{datetime_field}" column in ' + f'columns {columns}' ) - # Validate dataframe columns based on feature set entities + # Validate Avro columns based on feature set entities for feature_set in feature_sets_request: fs = self.get_feature_set( name=feature_set.name, version=feature_set.version @@ -431,10 +548,10 @@ def _validate_entity_rows_for_batch_retrieval( f"could not be found" ) for entity_type in fs.entities: - if entity_type.name not in entity_rows.columns: + if entity_type.name not in columns: raise ValueError( - f'Dataframe does not contain entity "{entity_type.name}"' - f' column in columns "{entity_rows.columns}"' + f'Input does not contain entity' + f' "{entity_type.name}" column in columns "{columns}"' ) def get_online_features( @@ -470,35 +587,55 @@ def get_online_features( ) # type: GetOnlineFeaturesResponse def ingest( - self, - feature_set: Union[str, FeatureSet], - source: Union[pd.DataFrame, str], - version: int = None, - force_update: bool = False, - max_workers: int = CPU_COUNT, - disable_progress_bar: bool = False, - chunk_size: int = 5000, - timeout: int = None, - ): + self, + feature_set: Union[str, FeatureSet], + source: Union[pd.DataFrame, str], + chunk_size: int = 10000, + version: int = None, + force_update: bool = False, + max_workers: int = max(CPU_COUNT - 1, 1), + disable_progress_bar: bool = False, + timeout: int = KAFKA_CHUNK_PRODUCTION_TIMEOUT + ) -> None: """ Loads feature data into Feast for a specific feature set. Args: - feature_set: Name of feature set or a feature set object - source: Either a file path or Pandas Dataframe to ingest into Feast + feature_set (typing.Union[str, FeatureSet]): + Feature set object or the string name of the feature set + (without a version). + + source (typing.Union[pd.DataFrame, str]): + Either a file path or Pandas Dataframe to ingest into Feast Files that are currently supported: - * parquet - * csv - * json - version: Feature set version - force_update: Automatically update feature set based on source data - prior to ingesting. This will also register changes to Feast - max_workers: Number of worker processes to use to encode values - disable_progress_bar: Disable printing of progress statistics - chunk_size: Maximum amount of rows to load into memory and ingest at - a time - timeout: Seconds to wait before ingestion times out + * parquet + * csv + * json + + chunk_size (int): + Amount of rows to load and ingest at a time. + + version (int): + Feature set version. + + force_update (bool): + Automatically update feature set based on source data prior to + ingesting. This will also register changes to Feast. + + max_workers (int): + Number of worker processes to use to encode values. + + disable_progress_bar (bool): + Disable printing of progress statistics. + + timeout (int): + Timeout in seconds to wait for completion. + + Returns: + None: + None """ + if isinstance(feature_set, FeatureSet): name = feature_set.name if version is None: @@ -508,38 +645,74 @@ def ingest( else: raise Exception(f"Feature set name must be provided") - table = _read_table_from_source(source) + # Read table and get row count + tmp_table_name = _read_table_from_source( + source, chunk_size, max_workers + ) + + pq_file = pq.ParquetFile(tmp_table_name) - # Update the feature set based on DataFrame schema - if force_update: - # Use a small as reference DataFrame to infer fields - ref_df = table.to_batches(max_chunksize=20)[0].to_pandas() + row_count = pq_file.metadata.num_rows - feature_set.infer_fields_from_df( - ref_df, discard_unused_fields=True, replace_existing_features=True + # Update the feature set based on PyArrow table of first row group + if force_update: + feature_set.infer_fields_from_pa( + table=pq_file.read_row_group(0), + discard_unused_fields=True, + replace_existing_features=True ) self.apply(feature_set) feature_set = self.get_feature_set(name, version) - if feature_set.source.source_type == "Kafka": - ingest_table_to_kafka( - feature_set=feature_set, - table=table, - max_workers=max_workers, - disable_pbar=disable_progress_bar, - chunk_size=chunk_size, - timeout=timeout, - ) - else: - raise Exception( - f"Could not determine source type for feature set " - f'"{feature_set.name}" with source type ' - f'"{feature_set.source.source_type}"' - ) + try: + # Kafka configs + brokers = feature_set.get_kafka_source_brokers() + topic = feature_set.get_kafka_source_topic() + producer = get_producer(brokers, row_count, disable_progress_bar) + + # Loop optimization declarations + produce = producer.produce + flush = producer.flush + + # Transform and push data to Kafka + if feature_set.source.source_type == "Kafka": + for chunk in get_feature_row_chunks( + file=tmp_table_name, + row_groups=list(range(pq_file.num_row_groups)), + fs=feature_set, + max_workers=max_workers): + + # Push FeatureRow one chunk at a time to kafka + for serialized_row in chunk: + produce(topic=topic, value=serialized_row) + + # Force a flush after each chunk + flush(timeout=timeout) + + # Remove chunk from memory + del chunk + + else: + raise Exception( + f"Could not determine source type for feature set " + f'"{feature_set.name}" with source type ' + f'"{feature_set.source.source_type}"' + ) + + # Print ingestion statistics + producer.print_results() + finally: + # Remove parquet file(s) that were created earlier + print("Removing temporary file(s)...") + os.remove(tmp_table_name) + + return None -def _build_feature_set_request(feature_ids: List[str]) -> List[FeatureSetRequest]: +def _build_feature_set_request( + feature_ids: List[str] +) -> List[FeatureSetRequest]: """ Builds a list of FeatureSet objects from feature set ids in order to retrieve feature data from Feast Serving @@ -566,18 +739,38 @@ def _build_feature_set_request(feature_ids: List[str]) -> List[FeatureSetRequest return list(feature_set_request.values()) -def _read_table_from_source(source: Union[pd.DataFrame, str]) -> pa.lib.Table: +def _read_table_from_source( + source: Union[pd.DataFrame, str], + chunk_size: int, + max_workers: int +) -> str: """ - Infers a data source type (path or Pandas Dataframe) and reads it in as + Infers a data source type (path or Pandas DataFrame) and reads it in as a PyArrow Table. + The PyArrow Table that is read will be written to a parquet file with row + group size determined by the minimum of: + * (table.num_rows / max_workers) + * chunk_size + + The parquet file that is created will be passed as file path to the + multiprocessing pool workers. + Args: - source: Either a string path or Pandas Dataframe + source (Union[pd.DataFrame, str]): + Either a string path or Pandas DataFrame. + + chunk_size (int): + Number of worker processes to use to encode values. + + max_workers (int): + Amount of rows to load and ingest at a time. Returns: - PyArrow table + str: Path to parquet file that was created. """ - # Pandas dataframe detected + + # Pandas DataFrame detected if isinstance(source, pd.DataFrame): table = pa.Table.from_pandas(df=source) @@ -597,8 +790,19 @@ def _read_table_from_source(source: Union[pd.DataFrame, str]) -> pa.lib.Table: else: table = pq.read_table(file_path) else: - raise ValueError(f"Unknown data source provided for ingestion: {source}") + raise ValueError( + f"Unknown data source provided for ingestion: {source}") # Ensure that PyArrow table is initialised assert isinstance(table, pa.lib.Table) - return table + + # Write table as parquet file with a specified row_group_size + tmp_table_name = f"{int(time.time())}.parquet" + row_group_size = min(ceil(table.num_rows / max_workers), chunk_size) + pq.write_table(table=table, where=tmp_table_name, + row_group_size=row_group_size) + + # Remove table from memory + del table + + return tmp_table_name diff --git a/sdk/python/feast/feature_set.py b/sdk/python/feast/feature_set.py index 893378e8fac..42979ea9116 100644 --- a/sdk/python/feast/feature_set.py +++ b/sdk/python/feast/feature_set.py @@ -13,21 +13,25 @@ # limitations under the License. -import pandas as pd -from typing import List, Optional from collections import OrderedDict from typing import Dict -from feast.source import Source -from pandas.api.types import is_datetime64_ns_dtype +from typing import List, Optional + +import pandas as pd +import pyarrow as pa +from feast.core.FeatureSet_pb2 import FeatureSetSpec as FeatureSetSpecProto from feast.entity import Entity from feast.feature import Feature, Field -from feast.core.FeatureSet_pb2 import FeatureSetSpec as FeatureSetSpecProto -from google.protobuf.duration_pb2 import Duration +from feast.loaders import yaml as feast_yaml +from feast.source import Source +from feast.type_map import DATETIME_COLUMN +from feast.type_map import pa_to_feast_value_type from feast.type_map import python_type_to_feast_value_type -from google.protobuf.json_format import MessageToJson from google.protobuf import json_format -from feast.type_map import DATETIME_COLUMN -from feast.loaders import yaml as feast_yaml +from google.protobuf.duration_pb2 import Duration +from google.protobuf.json_format import MessageToJson +from pandas.api.types import is_datetime64_ns_dtype +from pyarrow.lib import TimestampType class FeatureSet: @@ -256,7 +260,6 @@ def infer_fields_from_df( rows_to_sample: int = 100, ): """ - Adds fields (Features or Entities) to a feature set based on the schema of a Datatframe. Only Pandas dataframes are supported. All columns are detected as features, so setting at least one entity manually is @@ -283,6 +286,7 @@ def infer_fields_from_df( must have consistent types, even values within list types must be homogeneous """ + if entities is None: entities = list() if features is None: @@ -373,7 +377,187 @@ def infer_fields_from_df( self._fields = new_fields print(output_log) - def update_from_feature_set(self, feature_set): + def infer_fields_from_pa( + self, table: pa.lib.Table, + entities: Optional[List[Entity]] = None, + features: Optional[List[Feature]] = None, + replace_existing_features: bool = False, + replace_existing_entities: bool = False, + discard_unused_fields: bool = False + ) -> None: + """ + Adds fields (Features or Entities) to a feature set based on the schema + of a PyArrow table. Only PyArrow tables are supported. All columns are + detected as features, so setting at least one entity manually is + advised. + + + Args: + table (pyarrow.lib.Table): + PyArrow table to read schema from. + + entities (Optional[List[Entity]]): + List of entities that will be set manually and not inferred. + These will take precedence over any existing entities or + entities found in the PyArrow table. + + features (Optional[List[Feature]]): + List of features that will be set manually and not inferred. + These will take precedence over any existing feature or features + found in the PyArrow table. + + replace_existing_features (bool): + Boolean flag. If true, will replace existing features in this + feature set with features found in dataframe. If false, will + skip conflicting features. + + replace_existing_entities (bool): + Boolean flag. If true, will replace existing entities in this + feature set with features found in dataframe. If false, will + skip conflicting entities. + + discard_unused_fields (bool): + Boolean flag. Setting this to True will discard any existing + fields that are not found in the dataset or provided by the + user. + + Returns: + None: + None + """ + if entities is None: + entities = list() + if features is None: + features = list() + + # Validate whether the datetime column exists with the right name + if DATETIME_COLUMN not in table.column_names: + raise Exception("No column 'datetime'") + + # Validate the date type for the datetime column + if not isinstance(table.column(DATETIME_COLUMN).type, TimestampType): + raise Exception( + "Column 'datetime' does not have the correct type: datetime64[ms]" + ) + + # Create dictionary of fields that will not be inferred (manually set) + provided_fields = OrderedDict() + + for field in entities + features: + if not isinstance(field, Field): + raise Exception(f"Invalid field object type provided {type(field)}") + if field.name not in provided_fields: + provided_fields[field.name] = field + else: + raise Exception(f"Duplicate field name detected {field.name}.") + + new_fields = self._fields.copy() + output_log = "" + + # Add in provided fields + for name, field in provided_fields.items(): + if name in new_fields.keys(): + upsert_message = "created" + else: + upsert_message = "updated (replacing an existing field)" + + output_log += ( + f"{type(field).__name__} {field.name}" + f"({field.dtype}) manually {upsert_message}.\n" + ) + new_fields[name] = field + + # Iterate over all of the column names and create features + for column in table.column_names: + column = column.strip() + + # Skip datetime column + if DATETIME_COLUMN in column: + continue + + # Skip user provided fields + if column in provided_fields.keys(): + continue + + # Only overwrite conflicting fields if replacement is allowed + if column in new_fields: + if ( + isinstance(self._fields[column], Feature) + and not replace_existing_features + ): + continue + + if ( + isinstance(self._fields[column], Entity) + and not replace_existing_entities + ): + continue + + # Store this fields as a feature + # TODO: (Minor) Change the parameter name from dtype to patype + new_fields[column] = Feature( + name=column, + dtype=self._infer_pa_column_type(table.column(column)) + ) + + output_log += f"{type(new_fields[column]).__name__} {new_fields[column].name} ({new_fields[column].dtype}) added from PyArrow Table.\n" + + # Discard unused fields from feature set + if discard_unused_fields: + keys_to_remove = [] + for key in new_fields.keys(): + if not (key in table.column_names or key in provided_fields.keys()): + output_log += f"{type(new_fields[key]).__name__} {new_fields[key].name} ({new_fields[key].dtype}) removed because it is unused.\n" + keys_to_remove.append(key) + for key in keys_to_remove: + del new_fields[key] + + # Update feature set + self._fields = new_fields + print(output_log) + + def _infer_pd_column_type(self, column, series, rows_to_sample): + dtype = None + sample_count = 0 + + # Loop over all rows for this column to infer types + for key, value in series.iteritems(): + sample_count += 1 + # Stop sampling at the row limit + if sample_count > rows_to_sample: + continue + + # Infer the specific type for this row + current_dtype = python_type_to_feast_value_type(name=column, value=value) + + # Make sure the type is consistent for column + if dtype: + if dtype != current_dtype: + raise ValueError( + f"Type mismatch detected in column {column}. Both " + f"the types {current_dtype} and {dtype} " + f"have been found." + ) + else: + # Store dtype in field to type map if it isnt already + dtype = current_dtype + + return dtype + + def _infer_pa_column_type(self, column: pa.lib.ChunkedArray): + """ + Infers the PyArrow column type. + + :param column: Column from a PyArrow table + :type column: pa.lib.ChunkedArray + :return: + :rtype: + """ + # Validates the column to ensure that value types are consistent + column.validate() + return pa_to_feast_value_type(column) + + def _update_from_feature_set(self, feature_set): """ Deep replaces one feature set with another diff --git a/sdk/python/feast/job.py b/sdk/python/feast/job.py index 4273f86ea84..26f6181ee2d 100644 --- a/sdk/python/feast/job.py +++ b/sdk/python/feast/job.py @@ -1,12 +1,11 @@ import tempfile import time from datetime import datetime, timedelta -from typing import List +from typing import Iterable from urllib.parse import urlparse import fastavro import pandas as pd -from fastavro import reader as fastavro_reader from google.cloud import storage from feast.serving.ServingService_pb2 import GetJobRequest @@ -62,15 +61,18 @@ def reload(self): """ self.job_proto = self.serving_stub.GetJob(GetJobRequest(job=self.job_proto)).job - def result(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC): + def get_avro_files(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC): """ - Wait until job is done to get an iterable rows of result. - The row can only represent an Avro row in Feast 0.3. + Wait until job is done to get the file uri to Avro result files on + Google Cloud Storage. Args: - timeout_sec: max no of seconds to wait until job is done. If "timeout_sec" is exceeded, an exception will be raised. + timeout_sec (int): + Max no of seconds to wait until job is done. If "timeout_sec" + is exceeded, an exception will be raised. - Returns: Iterable of Avro rows + Returns: + str: Google Cloud Storage file uris of the returned Avro files. """ max_wait_datetime = datetime.now() + timedelta(seconds=timeout_sec) wait_duration_sec = 2 @@ -78,11 +80,13 @@ def result(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC): while self.status != JOB_STATUS_DONE: if datetime.now() > max_wait_datetime: raise Exception( - "Timeout exceeded while waiting for result. Please retry this method or use a longer timeout value." + "Timeout exceeded while waiting for result. Please retry " + "this method or use a longer timeout value." ) self.reload() time.sleep(wait_duration_sec) + # Backoff the wait duration exponentially up till MAX_WAIT_INTERVAL_SEC wait_duration_sec = min(wait_duration_sec * 2, MAX_WAIT_INTERVAL_SEC) @@ -95,7 +99,22 @@ def result(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC): "your Feast Serving deployment." ) - uris = [urlparse(uri) for uri in self.job_proto.file_uris] + return [urlparse(uri) for uri in self.job_proto.file_uris] + + def result(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC): + """ + Wait until job is done to get an iterable rows of result. The row can + only represent an Avro row in Feast 0.3. + + Args: + timeout_sec (int): + Max no of seconds to wait until job is done. If "timeout_sec" + is exceeded, an exception will be raised. + + Returns: + Iterable of Avro rows. + """ + uris = self.get_avro_files(timeout_sec) for file_uri in uris: if file_uri.scheme == "gs": file_obj = tempfile.TemporaryFile() @@ -113,16 +132,64 @@ def result(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC): for record in avro_reader: yield record - def to_dataframe(self, timeout_sec: int = DEFAULT_TIMEOUT_SEC): + def to_dataframe( + self, + timeout_sec: int = DEFAULT_TIMEOUT_SEC + ) -> pd.DataFrame: """ - Wait until job is done to get an interable rows of result + Wait until a job is done to get an iterable rows of result. This method + will split the response into chunked DataFrame of a specified size to + to be yielded to the instance calling it. Args: - timeout_sec: max no of seconds to wait until job is done. If "timeout_sec" is exceeded, an exception will be raised. - Returns: pandas Dataframe of the feature values + max_chunk_size (int): + Maximum number of rows that the DataFrame should contain. + + timeout_sec (int): + Max no of seconds to wait until job is done. If "timeout_sec" + is exceeded, an exception will be raised. + + Returns: + pd.DataFrame: + Pandas DataFrame of the feature values. """ records = [r for r in self.result(timeout_sec=timeout_sec)] return pd.DataFrame.from_records(records) + def to_chunked_dataframe( + self, + max_chunk_size: int = -1, + timeout_sec: int = DEFAULT_TIMEOUT_SEC + ) -> pd.DataFrame: + """ + Wait until a job is done to get an iterable rows of result. This method + will split the response into chunked DataFrame of a specified size to + to be yielded to the instance calling it. + + Args: + max_chunk_size (int): + Maximum number of rows that the DataFrame should contain. + + timeout_sec (int): + Max no of seconds to wait until job is done. If "timeout_sec" + is exceeded, an exception will be raised. + + Returns: + pd.DataFrame: + Pandas DataFrame of the feature values. + """ + # Max chunk size defined by user + records = [] + for result in self.result(timeout_sec=timeout_sec): + result.append(records) + if len(records) == max_chunk_size: + df = pd.DataFrame.from_records(records) + records.clear() # Empty records array + yield df + + # Handle for last chunk that is < max_chunk_size + if not records: + yield pd.DataFrame.from_records(records) + def __iter__(self): return iter(self.result()) diff --git a/sdk/python/feast/loaders/abstract_producer.py b/sdk/python/feast/loaders/abstract_producer.py new file mode 100644 index 00000000000..884ae49984c --- /dev/null +++ b/sdk/python/feast/loaders/abstract_producer.py @@ -0,0 +1,248 @@ +# Copyright 2019 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Union + +from tqdm import tqdm + + +class AbstractProducer: + """ + Abstract class for Kafka producers + """ + + def __init__( + self, + brokers: str, + row_count: int, + disable_progress_bar: bool + ): + self.brokers = brokers + self.row_count = row_count + self.error_count = 0 + self.last_exception = "" + + # Progress bar will always display average rate + self.pbar = tqdm( + total=row_count, + unit="rows", + smoothing=0, + disable=disable_progress_bar + ) + + def produce(self, topic: str, data: str): + message = "{} should implement a produce method".format( + self.__class__.__name__) + raise NotImplementedError(message) + + def flush(self, timeout: int): + message = "{} should implement a flush method".format( + self.__class__.__name__) + raise NotImplementedError(message) + + def _inc_pbar(self, meta): + self.pbar.update(1) + + def _set_error(self, exception: str): + self.error_count += 1 + self.last_exception = exception + + def print_results(self) -> None: + """ + Print ingestion statistics. + + Returns: + None: None + """ + # Refresh and close tqdm progress bar + self.pbar.refresh() + + self.pbar.close() + + print("Ingestion complete!") + + failed_message = ( + "" + if self.error_count == 0 + else f"\nFail: {self.error_count / self.row_count}" + ) + + last_exception_message = ( + "" + if self.last_exception == "" + else f"\nLast exception:\n{self.last_exception}" + ) + + print( + f"\nIngestion statistics:" + f"\nSuccess: {self.pbar.n}/{self.row_count}" + f"{failed_message}" + f"{last_exception_message}" + ) + return None + + +class ConfluentProducer(AbstractProducer): + """ + Concrete implementation of Confluent Kafka producer (confluent-kafka) + """ + + def __init__( + self, + brokers: str, + row_count: int, + disable_progress_bar: bool + ): + from confluent_kafka import Producer + self.producer = Producer({"bootstrap.servers": brokers}) + super().__init__(brokers, row_count, disable_progress_bar) + + def produce(self, topic: str, value: bytes) -> None: + """ + Generic produce that implements confluent-kafka's produce method to + push a byte encoded object into a Kafka topic. + + Args: + topic (str): Kafka topic. + value (bytes): Byte encoded object. + + Returns: + None: None. + """ + + try: + self.producer.produce( + topic, value=value, callback=self._delivery_callback) + # Serve delivery callback queue. + # NOTE: Since produce() is an asynchronous API this poll() call + # will most likely not serve the delivery callback for the + # last produce()d message. + self.producer.poll(0) + except Exception as ex: + self._set_error(str(ex)) + + return None + + def flush(self, timeout: Optional[int]): + """ + Generic flush that implements confluent-kafka's flush method. + + Args: + timeout (Optional[int]): Timeout in seconds to wait for completion. + + Returns: + int: Number of messages still in queue. + """ + return self.producer.flush(timeout=timeout) + + def _delivery_callback(self, err: str, msg) -> None: + """ + Optional per-message delivery callback (triggered by poll() or flush()) + when a message has been successfully delivered or permanently failed + delivery (after retries). + + Although the msg argument is not used, the current method signature is + required as specified in the confluent-kafka documentation. + + Args: + err (str): Error message. + msg (): Kafka message. + + Returns: + None + """ + if err: + self._set_error(err) + else: + self._inc_pbar(None) + + +class KafkaPythonProducer(AbstractProducer): + """ + Concrete implementation of Python Kafka producer (kafka-python) + """ + + def __init__( + self, + brokers: str, + row_count: int, + disable_progress_bar: bool + ): + from kafka import KafkaProducer + self.producer = KafkaProducer(bootstrap_servers=[brokers]) + super().__init__(brokers, row_count, disable_progress_bar) + + def produce(self, topic: str, value: bytes): + """ + Generic produce that implements kafka-python's send method to push a + byte encoded object into a Kafka topic. + + Args: + topic (str): Kafka topic. + value (bytes): Byte encoded object. + + Returns: + FutureRecordMetadata: resolves to RecordMetadata + + Raises: + KafkaTimeoutError: if unable to fetch topic metadata, or unable + to obtain memory buffer prior to configured max_block_ms + """ + return self.producer.send(topic, value=value).add_callback( + self._inc_pbar).add_errback(self._set_error) + + def flush(self, timeout: Optional[int]): + """ + Generic flush that implements kafka-python's flush method. + + Args: + timeout (Optional[int]): timeout in seconds to wait for completion. + + Returns: + None + + Raises: + KafkaTimeoutError: failure to flush buffered records within the + provided timeout + """ + return self.producer.flush(timeout=timeout) + + +def get_producer( + brokers: str, row_count: int, disable_progress_bar: bool +) -> Union[ConfluentProducer, KafkaPythonProducer]: + """ + Simple context helper function that returns a AbstractProducer object when + invoked. + + This helper function will try to import confluent-kafka as a producer first. + + This helper function will fallback to kafka-python if it fails to import + confluent-kafka. + + Args: + brokers (str): Kafka broker information with hostname and port. + row_count (int): Number of rows in table + + Returns: + Union[ConfluentProducer, KafkaPythonProducer]: + Concrete implementation of a Kafka producer. Ig can be: + * confluent-kafka producer + * kafka-python producer + """ + try: + return ConfluentProducer(brokers, row_count, disable_progress_bar) + except ImportError as e: + print("Unable to import confluent-kafka, falling back to kafka-python") + return KafkaPythonProducer(brokers, row_count, disable_progress_bar) diff --git a/sdk/python/feast/loaders/file.py b/sdk/python/feast/loaders/file.py index 8dd6b503a74..108f2790dd8 100644 --- a/sdk/python/feast/loaders/file.py +++ b/sdk/python/feast/loaders/file.py @@ -1,70 +1,159 @@ +# Copyright 2019 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import re import shutil import tempfile -from typing import Optional -from urllib.parse import urlparse import uuid -import pandas as pd from datetime import datetime +from typing import List, Optional, Tuple, Union +from urllib.parse import urlparse, ParseResult + +import pandas as pd from google.cloud import storage from pandavro import to_avro -def export_dataframe_to_staging_location( - df: pd.DataFrame, staging_location_uri: str -) -> str: +def export_source_to_staging_location( + source: Union[pd.DataFrame, str], staging_location_uri: str +) -> List[str]: """ - Uploads a dataframe to a remote staging location + Uploads a DataFrame as an Avro file to a remote staging location. + + The local staging location specified in this function is used for E2E + tests, please do not use it. Args: - df: Pandas dataframe - staging_location_uri: Remote staging location where dataframe should be written + source (Union[pd.DataFrame, str]: + Source of data to be staged. Can be a pandas DataFrame or a file + path. + + Only three types of source are allowed: + * Pandas DataFrame + * Local Avro file + * GCS Avro file + + + staging_location_uri (str): + Remote staging location where DataFrame should be written. Examples: - gs://bucket/path/ - file:///data/subfolder/ + * gs://bucket/path/ + * file:///data/subfolder/ Returns: - Returns the full path to the file in the remote staging location + List[str]: + Returns a list containing the full path to the file(s) in the + remote staging location. """ - # Validate staging location uri = urlparse(staging_location_uri) + + # Prepare Avro file to be exported to staging location + if isinstance(source, pd.DataFrame): + # DataFrame provided as a source + if uri.scheme == "file": + uri_path = uri.path + else: + uri_path = None + + # Remote gs staging location provided by serving + dir_path, file_name, source_path = export_dataframe_to_local( + source, + uri_path + ) + elif urlparse(source).scheme in ["", "file"]: + # Local file provided as a source + dir_path = None + file_name = os.path.basename(source) + source_path = os.path.abspath(os.path.join( + urlparse(source).netloc, urlparse(source).path)) + elif urlparse(source).scheme == "gs": + # Google Cloud Storage path provided + input_source_uri = urlparse(source) + if "*" in source: + # Wildcard path + return _get_files( + bucket=input_source_uri.hostname, + uri=input_source_uri + ) + else: + return [source] + else: + raise Exception(f"Only string and DataFrame types are allowed as a " + f"source, {type(source)} was provided.") + + # Push data to required staging location if uri.scheme == "gs": - dir_path, file_name, source_path = export_dataframe_to_local(df) + # Staging location is a Google Cloud Storage path upload_file_to_gcs( - source_path, uri.hostname, str(uri.path).strip("/") + "/" + file_name + source_path, + uri.hostname, + str(uri.path).strip("/") + "/" + file_name ) - if len(str(dir_path)) < 5: - raise Exception(f"Export location {dir_path} dangerous. Stopping.") - shutil.rmtree(dir_path) elif uri.scheme == "file": - dir_path, file_name, source_path = export_dataframe_to_local(df, uri.path) + # Staging location is a file path + # Used for end-to-end test + pass else: raise Exception( - f"Staging location {staging_location_uri} does not have a valid URI. Only gs:// and file:// are supported" + f"Staging location {staging_location_uri} does not have a " + f"valid URI. Only gs:// and file:// uri scheme are supported." ) - return staging_location_uri.rstrip("/") + "/" + file_name + # Clean up, remove local staging file + if isinstance(source, pd.DataFrame) and len(str(dir_path)) > 4: + shutil.rmtree(dir_path) + + return [staging_location_uri.rstrip("/") + "/" + file_name] -def export_dataframe_to_local(df: pd.DataFrame, dir_path: Optional[str] = None): +def export_dataframe_to_local( + df: pd.DataFrame, + dir_path: Optional[str] = None +) -> Tuple[str, str, str]: """ - Exports a pandas dataframe to the local filesystem + Exports a pandas DataFrame to the local filesystem. Args: - df: Pandas dataframe to save - dir_path: (optional) Absolute directory path '/data/project/subfolder/' + df (pd.DataFrame): + Pandas DataFrame to save. + + dir_path (Optional[str]): + Absolute directory path '/data/project/subfolder/'. + + Returns: + Tuple[str, str, str]: + Tuple of directory path, file name and destination path. The + destination path can be obtained by concatenating the directory + path and file name. """ # Create local staging location if not provided if dir_path is None: dir_path = tempfile.mkdtemp() - file_name = f'{datetime.now().strftime("%d-%m-%Y_%I-%M-%S_%p")}_{str(uuid.uuid4())[:8]}.avro' + file_name = _get_file_name() dest_path = f"{dir_path}/{file_name}" # Temporarily rename datetime column to event_timestamp. Ideally we would # force the schema with our avro writer instead. - df.columns = ["event_timestamp" if col == "datetime" else col for col in df.columns] + df.columns = [ + "event_timestamp" + if col == "datetime" else col + for col in df.columns + ] try: # Export dataset to file in local path @@ -74,23 +163,84 @@ def export_dataframe_to_local(df: pd.DataFrame, dir_path: Optional[str] = None): finally: # Revert event_timestamp column to datetime df.columns = [ - "datetime" if col == "event_timestamp" else col for col in df.columns + "datetime" + if col == "event_timestamp" else col + for col in df.columns ] return dir_path, file_name, dest_path -def upload_file_to_gcs(local_path: str, bucket: str, remote_path: str): +def upload_file_to_gcs(local_path: str, bucket: str, remote_path: str) -> None: """ - Upload a file from the local file system to Google Cloud Storage (GCS) + Upload a file from the local file system to Google Cloud Storage (GCS). Args: - local_path: Local filesystem path of file to upload - bucket: GCS bucket to upload to - remote_path: Path within GCS bucket to upload file to, includes file name + local_path (str): + Local filesystem path of file to upload. + + bucket (str): + GCS bucket destination to upload to. + + remote_path (str): + Path within GCS bucket to upload file to, includes file name. + + Returns: + None: + None """ storage_client = storage.Client(project=None) bucket = storage_client.get_bucket(bucket) blob = bucket.blob(remote_path) blob.upload_from_filename(local_path) + + +def _get_files(bucket: str, uri: ParseResult) -> List[str]: + """ + List all available files within a Google storage bucket that matches a wild + card path. + + Args: + bucket (str): + Google Storage bucket to reference. + + uri (urllib.parse.ParseResult): + Wild card uri path containing the "*" character. + Example: + * gs://feast/staging_location/* + * gs://feast/staging_location/file_*.avro + + Returns: + List[str]: + List of all available files matching the wildcard path. + """ + + storage_client = storage.Client(project=None) + bucket = storage_client.get_bucket(bucket) + path = uri.path + + if "*" in path: + regex = re.compile(path.replace("*", ".*?").strip("/")) + blob_list = bucket.list_blobs( + prefix=path.strip("/").split("*")[0], + delimiter="/" + ) + # File path should not be in path (file path must be longer than path) + return [f"{uri.scheme}://{uri.hostname}/{file}" + for file in [x.name for x in blob_list] + if re.match(regex, file) and file not in path] + else: + raise Exception(f"{path} is not a wildcard path") + + +def _get_file_name() -> str: + """ + Create a random file name. + + Returns: + str: + Randomised file name. + """ + + return f'{datetime.now().strftime("%d-%m-%Y_%I-%M-%S_%p")}_{str(uuid.uuid4())[:8]}.avro' diff --git a/sdk/python/feast/loaders/ingest.py b/sdk/python/feast/loaders/ingest.py index 23ba2ecb3b4..527ab481fe0 100644 --- a/sdk/python/feast/loaders/ingest.py +++ b/sdk/python/feast/loaders/ingest.py @@ -1,18 +1,16 @@ import logging -import multiprocessing -import os -import time from functools import partial -from multiprocessing import Process, Queue, Pool -from typing import Iterable +from multiprocessing import Pool +from typing import Iterable, List + import pandas as pd -import pyarrow as pa +import pyarrow.parquet as pq +from feast.constants import DATETIME_COLUMN from feast.feature_set import FeatureSet -from feast.type_map import convert_dict_to_proto_values +from feast.type_map import pa_column_to_timestamp_proto_column, \ + pa_column_to_proto_column +from feast.types import Field_pb2 as FieldProto from feast.types.FeatureRow_pb2 import FeatureRow -from kafka import KafkaProducer -from tqdm import tqdm -from feast.constants import DATETIME_COLUMN _logger = logging.getLogger(__name__) @@ -21,221 +19,120 @@ FEAST_SERVING_URL_ENV_KEY = "FEAST_SERVING_URL" # type: str FEAST_CORE_URL_ENV_KEY = "FEAST_CORE_URL" # type: str BATCH_FEATURE_REQUEST_WAIT_TIME_SECONDS = 300 -CPU_COUNT = os.cpu_count() # type: int KAFKA_CHUNK_PRODUCTION_TIMEOUT = 120 # type: int -def _kafka_feature_row_producer( - feature_row_queue: Queue, row_count: int, brokers, topic, ctx: dict, pbar: tqdm -): +def _encode_pa_tables( + file: str, + fs: FeatureSet, + row_group_idx: int, +) -> List[bytes]: """ - Pushes Feature Rows to Kafka. Reads rows from a queue. Function will run - until total row_count is reached. + Helper function to encode a PyArrow table(s) read from parquet file(s) into + FeatureRows. - Args: - feature_row_queue: Queue containing feature rows. - row_count: Total row count to process - brokers: Broker to push to - topic: Topic to push to - ctx: Context dict used to communicate with primary process - pbar: Progress bar object - """ - - # Callback for failed production to Kafka - def on_error(e): - # Save last exception - ctx["last_exception"] = e - - # Increment error count - if "error_count" in ctx: - ctx["error_count"] += 1 - else: - ctx["error_count"] = 1 - - # Callback for succeeded production to Kafka - def on_success(meta): - pbar.update() - - producer = KafkaProducer(bootstrap_servers=brokers) - processed_rows = 0 - - # Loop through feature rows until all rows are processed - while processed_rows < row_count: - # Wait if queue is empty - if feature_row_queue.empty(): - time.sleep(1) - producer.flush(timeout=KAFKA_CHUNK_PRODUCTION_TIMEOUT) - else: - while not feature_row_queue.empty(): - row = feature_row_queue.get() - if row is not None: - # Push row to Kafka - producer.send(topic, row.SerializeToString()).add_callback( - on_success - ).add_errback(on_error) - processed_rows += 1 - - # Force an occasional flush - if processed_rows % 10000 == 0: - producer.flush(timeout=KAFKA_CHUNK_PRODUCTION_TIMEOUT) - del row - pbar.refresh() - - # Ensure that all rows are pushed - producer.flush(timeout=KAFKA_CHUNK_PRODUCTION_TIMEOUT) - - # Using progress bar as counter is much faster than incrementing dict - ctx["success_count"] = pbar.n - pbar.close() - - -def _encode_pa_chunks( - tbl: pa.lib.Table, - fs: FeatureSet, - max_workers: int, - df_datetime_dtype: pd.DataFrame.dtypes, - chunk_size: int = 5000, -) -> Iterable[FeatureRow]: - """ - Generator function to encode rows in PyArrow table to FeatureRows by - breaking up the table into batches. + This function accepts a list of file directory pointing to many parquet + files. All parquet files must have the same schema. - Each batch will have its rows spread accross a pool of workers to be - transformed into FeatureRow objects. + Each parquet file will be read into as a table and encoded into FeatureRows + using a pool of max_workers workers. Args: - tbl: PyArrow table to be processed. - fs: FeatureSet describing PyArrow table. - max_workers: Maximum number of workers. - df_datetime_dtype: Pandas dtype of datetime column. - chunk_size: Maximum size of each chunk when PyArrow table is batched. - - Returns: - Iterable FeatureRow object. - """ - - pool = Pool(max_workers) - - # Create a partial function with static non-iterable arguments - func = partial( - convert_dict_to_proto_values, - df_datetime_dtype=df_datetime_dtype, - feature_set=fs, - ) - - for batch in tbl.to_batches(max_chunksize=chunk_size): - m_df = batch.to_pandas() - results = pool.map_async(func, m_df.to_dict("records")) - yield from results.get() + file (str): + File directory of all the parquet file to encode. + Parquet file must have more than one row group. - pool.close() - pool.join() - return + fs (feast.feature_set.FeatureSet): + FeatureSet describing parquet files. + row_group_idx(int): + Row group index to read and encode into byte like FeatureRow + protobuf objects. -def ingest_table_to_kafka( - feature_set: FeatureSet, - table: pa.lib.Table, - max_workers: int, - chunk_size: int = 5000, - disable_pbar: bool = False, - timeout: int = None, -) -> None: + Returns: + List[bytes]: + List of byte encoded FeatureRows from the parquet file. """ - Ingest a PyArrow Table to a Kafka topic based for a Feature Set + pq_file = pq.ParquetFile(file) + # Read parquet file as a PyArrow table + table = pq_file.read_row_group(row_group_idx) + + # Add datetime column + datetime_col = pa_column_to_timestamp_proto_column( + table.column(DATETIME_COLUMN)) + + # Preprocess the columns by converting all its values to Proto values + proto_columns = { + field_name: pa_column_to_proto_column(field.dtype, + table.column(field_name)) + for field_name, field in fs.fields.items() + } + + feature_set = f"{fs.name}:{fs.version}" + + # List to store result + feature_rows = [] + + # Loop optimization declaration(s) + field = FieldProto.Field + proto_items = proto_columns.items() + append = feature_rows.append + + # Iterate through the rows + for row_idx in range(table.num_rows): + feature_row = FeatureRow(event_timestamp=datetime_col[row_idx], + feature_set=feature_set) + # Loop optimization declaration + ext = feature_row.fields.extend + + # Insert field from each column + for k, v in proto_items: + ext([field(name=k, value=v[row_idx])]) + + # Append FeatureRow in byte string form + append(feature_row.SerializeToString()) + + return feature_rows + + +def get_feature_row_chunks( + file: str, + row_groups: List[int], + fs: FeatureSet, + max_workers: int +) -> Iterable[List[bytes]]: + """ + Iterator function to encode a PyArrow table read from a parquet file to + FeatureRow(s). Args: - feature_set: FeatureSet describing PyArrow table. - table: PyArrow table to be processed. - max_workers: Maximum number of workers. - chunk_size: Maximum size of each chunk when PyArrow table is batched. - disable_pbar: Flag to indicate if tqdm progress bar should be disabled. - timeout: Maximum time before method times out - """ + file (str): + File directory of the parquet file. The parquet file must have more + than one row group. - pbar = tqdm(unit="rows", total=table.num_rows, disable=disable_pbar) - - # Use a small DataFrame to validate feature set schema - ref_df = table.to_batches(max_chunksize=100)[0].to_pandas() - df_datetime_dtype = ref_df[DATETIME_COLUMN].dtype - - # Validate feature set schema - _validate_dataframe(ref_df, feature_set) - - # Create queue through which encoding and production will coordinate - row_queue = Queue() - - # Create a context object to send and receive information across processes - ctx = multiprocessing.Manager().dict( - {"success_count": 0, "error_count": 0, "last_exception": ""} - ) - - # Create producer to push feature rows to Kafka - ingestion_process = Process( - target=_kafka_feature_row_producer, - args=( - row_queue, - table.num_rows, - feature_set.get_kafka_source_brokers(), - feature_set.get_kafka_source_topic(), - ctx, - pbar, - ), - ) - - try: - # Start ingestion process - print( - f"\n(ingest table to kafka) Ingestion started for {feature_set.name}:{feature_set.version}" - ) - ingestion_process.start() - - # Iterate over chunks in the table and return feature rows - for row in _encode_pa_chunks( - tbl=table, - fs=feature_set, - max_workers=max_workers, - chunk_size=chunk_size, - df_datetime_dtype=df_datetime_dtype, - ): - # Push rows onto a queue for the production process to pick up - row_queue.put(row) - while row_queue.qsize() > chunk_size: - time.sleep(0.1) - row_queue.put(None) - except Exception as ex: - _logger.error(f"Exception occurred: {ex}") - finally: - # Wait for the Kafka production to complete - ingestion_process.join(timeout=timeout) - failed_message = ( - "" - if ctx["error_count"] == 0 - else f"\nFail: {ctx['error_count']}/{table.num_rows}" - ) + row_groups (List[int]): + Specific row group indexes to be read and transformed in the parquet + file. - last_exception_message = ( - "" - if ctx["last_exception"] == "" - else f"\nLast exception:\n{ctx['last_exception']}" - ) - print( - f"\nIngestion statistics:" - f"\nSuccess: {ctx['success_count']}/{table.num_rows}" - f"{failed_message}" - f"{last_exception_message}" - ) + fs (feast.feature_set.FeatureSet): + FeatureSet describing parquet files. + max_workers (int): + Maximum number of workers to spawn. -def _validate_dataframe(dataframe: pd.DataFrame, feature_set: FeatureSet): + Returns: + Iterable[List[bytes]]: + Iterable list of byte encoded FeatureRow(s). """ - Validates a Pandas dataframe based on a feature set - Args: - dataframe: Pandas dataframe - feature_set: Feature Set instance - """ + pool = Pool(max_workers) + func = partial(_encode_pa_tables, file, fs) + for chunk in pool.imap_unordered(func, row_groups): + yield chunk + return + +def validate_dataframe(dataframe: pd.DataFrame, feature_set: FeatureSet): if "datetime" not in dataframe.columns: raise ValueError( f'Dataframe does not contain entity "datetime" in columns {dataframe.columns}' diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py index 7573276d74a..ca13c2573bc 100644 --- a/sdk/python/feast/type_map.py +++ b/sdk/python/feast/type_map.py @@ -12,12 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from datetime import datetime, timezone +from typing import List + import numpy as np import pandas as pd -from datetime import datetime, timezone -from feast.value_type import ValueType +import pyarrow as pa +from feast.constants import DATETIME_COLUMN +from feast.types import ( + FeatureRow_pb2 as FeatureRowProto, + Field_pb2 as FieldProto, +) from feast.types.Value_pb2 import ( Value as ProtoValue, + ValueType as ProtoValueType, Int64List, Int32List, BoolList, @@ -26,9 +34,9 @@ StringList, FloatList, ) -from feast.types import FeatureRow_pb2 as FeatureRowProto, Field_pb2 as FieldProto +from feast.value_type import ValueType from google.protobuf.timestamp_pb2 import Timestamp -from feast.constants import DATETIME_COLUMN +from pyarrow.lib import TimestampType def python_type_to_feast_value_type( @@ -104,9 +112,9 @@ def python_type_to_feast_value_type( return ValueType[common_item_value_type.name + "_LIST"] else: raise ValueError( - f"Value type for field {name} is {value.dtype.__str__()} " - f"but recursion is not allowed. Array types can only be one " - f"level deep." + f"Value type for field {name} is {value.dtype.__str__()} but " + f"recursion is not allowed. Array types can only be one level " + f"deep." ) return type_map[value.dtype.__str__()] @@ -160,7 +168,7 @@ def convert_series_to_proto_values(row: pd.Series): def convert_dict_to_proto_values( - row: dict, df_datetime_dtype: pd.DataFrame.dtypes, feature_set + row: dict, df_datetime_dtype: pd.DataFrame.dtypes, feature_set ) -> FeatureRowProto.FeatureRow: """ Encode a dictionary describing a feature row into a FeatureRows object. @@ -211,12 +219,14 @@ def _pd_datetime_to_timestamp_proto(dtype, value) -> Timestamp: # If timestamp does not contain timezone, we assume it is of local # timezone and adjust it to UTC local_timezone = datetime.now(timezone.utc).astimezone().tzinfo - value = value.tz_localize(local_timezone).tz_convert("UTC").tz_localize(None) + value = value.tz_localize(local_timezone).tz_convert("UTC").tz_localize( + None) return Timestamp(seconds=int(value.timestamp())) if dtype.__str__() == "datetime64[ns, UTC]": return Timestamp(seconds=int(value.timestamp())) else: - return Timestamp(seconds=np.datetime64(value).astype("int64") // 1000000) + return Timestamp( + seconds=np.datetime64(value).astype("int64") // 1000000) def _type_err(item, dtype): @@ -344,3 +354,139 @@ def _python_value_to_proto_value(feast_value_type, value) -> ProtoValue: return ProtoValue(bool_val=value) raise Exception(f"Unsupported data type: ${str(type(value))}") + +def pa_to_feast_value_attr(pa_type: object): + """ + Returns the equivalent Feast ValueType string for the given pa.lib type. + + Args: + pa_type (object): + PyArrow type. + + Returns: + str: + Feast attribute name in Feast ValueType string-ed representation. + """ + # Mapping of PyArrow type to attribute name in Feast ValueType strings + type_map = { + "timestamp[ms]": "int64_val", + "int32": "int32_val", + "int64": "int64_val", + "double": "double_val", + "float": "float_val", + "string": "string_val", + "binary": "bytes_val", + "bool": "bool_val", + "list": "int32_list_val", + "list": "int64_list_val", + "list": "double_list_val", + "list": "float_list_val", + "list": "string_list_val", + "list": "bytes_list_val", + "list": "bool_list_val", + } + + return type_map[pa_type.__str__()] + + +def pa_to_value_type(pa_type: object): + """ + Returns the equivalent Feast ValueType for the given pa.lib type. + + Args: + pa_type (object): + PyArrow type. + + Returns: + feast.types.Value_pb2.ValueType: + Feast ValueType. + + """ + + # Mapping of PyArrow to attribute name in Feast ValueType + type_map = { + "timestamp[ms]": ProtoValueType.INT64, + "int32": ProtoValueType.INT32, + "int64": ProtoValueType.INT64, + "double": ProtoValueType.DOUBLE, + "float": ProtoValueType.FLOAT, + "string": ProtoValueType.STRING, + "binary": ProtoValueType.BYTES, + "bool": ProtoValueType.BOOL, + "list": ProtoValueType.INT32_LIST, + "list": ProtoValueType.INT64_LIST, + "list": ProtoValueType.DOUBLE_LIST, + "list": ProtoValueType.FLOAT_LIST, + "list": ProtoValueType.STRING_LIST, + "list": ProtoValueType.BYTES_LIST, + "list": ProtoValueType.BOOL_LIST, + } + return type_map[pa_type.__str__()] + + +def pa_to_feast_value_type( + value: object +) -> ValueType: + type_map = { + "timestamp[ms]": ValueType.INT64, + "int32": ValueType.INT32, + "int64": ValueType.INT64, + "double": ValueType.DOUBLE, + "float": ValueType.FLOAT, + "string": ValueType.STRING, + "binary": ValueType.BYTES, + "bool": ValueType.BOOL, + "list": ValueType.INT32_LIST, + "list": ValueType.INT64_LIST, + "list": ValueType.DOUBLE_LIST, + "list": ValueType.FLOAT_LIST, + "list": ValueType.STRING_LIST, + "list": ValueType.BYTES_LIST, + "list": ValueType.BOOL_LIST, + } + return type_map[value.type.__str__()] + + +def pa_column_to_timestamp_proto_column( + column: pa.lib.ChunkedArray +) -> Timestamp: + if not isinstance(column.type, TimestampType): + raise Exception("Only TimestampType columns are allowed") + + proto_column = [] + for val in column: + timestamp = Timestamp() + timestamp.FromMicroseconds( + micros=int(val.as_py().timestamp() * 1_000_000)) + proto_column.append(timestamp) + return proto_column + + +def pa_column_to_proto_column( + feast_value_type, + column: pa.lib.ChunkedArray +) -> List[ProtoValue]: + type_map = {ValueType.INT32: "int32_val", + ValueType.INT64: "int64_val", + ValueType.FLOAT: "float_val", + ValueType.DOUBLE: "double_val", + ValueType.STRING: "string_val", + ValueType.BYTES: "bytes_val", + ValueType.BOOL: "bool_val", + ValueType.BOOL_LIST: {"bool_list_val": BoolList}, + ValueType.BYTES_LIST: {"bytes_list_val": BytesList}, + ValueType.STRING_LIST: {"string_list_val": StringList}, + ValueType.FLOAT_LIST: {"float_list_val": FloatList}, + ValueType.DOUBLE_LIST: {"double_list_val": DoubleList}, + ValueType.INT32_LIST: {"int32_list_val": Int32List}, + ValueType.INT64_LIST: {"int64_list_val": Int64List}, } + + value = type_map[feast_value_type] + # Process list types + if type(value) == dict: + list_param_name = list(value.keys())[0] + return [ProtoValue( + **{list_param_name: value[list_param_name](val=x.as_py())}) + for x in column] + else: + return [ProtoValue(**{value: x.as_py()}) for x in column] diff --git a/sdk/python/setup.py b/sdk/python/setup.py index 66cad904b01..420d1a39afc 100644 --- a/sdk/python/setup.py +++ b/sdk/python/setup.py @@ -32,7 +32,7 @@ "googleapis-common-protos==1.*", "google-cloud-bigquery-storage==0.7.*", "grpcio==1.*", - "pandas==0.*", + "pandas>=0.25.0", "pandavro==1.5.*", "protobuf>=3.10", "PyYAML==5.1.*", @@ -44,6 +44,7 @@ "pyarrow>=0.15.1", "numpy", "google", + "confluent_kafka" ] # README file from Feast repo root directory diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py index 2243ebfd1b3..9ef6e3d56fb 100644 --- a/sdk/python/tests/test_client.py +++ b/sdk/python/tests/test_client.py @@ -381,7 +381,7 @@ def test_feature_set_ingest_success(self, dataframe, client, mocker): ) # Need to create a mock producer - with patch("feast.loaders.ingest.KafkaProducer") as mocked_queue: + with patch("feast.client.get_producer") as mocked_queue: # Ingest data into Feast client.ingest("driver-feature-set", dataframe) @@ -445,6 +445,6 @@ def test_feature_set_types_success(self, client, dataframe, mocker): ) # Need to create a mock producer - with patch("feast.loaders.ingest.KafkaProducer") as mocked_queue: + with patch("feast.client.get_producer") as mocked_queue: # Ingest data into Feast client.ingest(all_types_fs, dataframe) diff --git a/tests/e2e/bq-batch-retrieval.py b/tests/e2e/bq-batch-retrieval.py index 2d6668eaa86..6b654f9a70e 100644 --- a/tests/e2e/bq-batch-retrieval.py +++ b/tests/e2e/bq-batch-retrieval.py @@ -2,6 +2,7 @@ import time from datetime import datetime from datetime import timedelta +from urllib.parse import urlparse import numpy as np import pandas as pd @@ -12,7 +13,9 @@ from feast.feature import Feature from feast.feature_set import FeatureSet from feast.type_map import ValueType +from google.cloud import storage from google.protobuf.duration_pb2 import Duration +from pandavro import to_avro pd.set_option('display.max_columns', None) @@ -31,6 +34,11 @@ def allow_dirty(pytestconfig): return True if pytestconfig.getoption("allow_dirty").lower() == "true" else False +@pytest.fixture(scope="module") +def gcs_path(pytestconfig): + return pytestconfig.getoption("gcs_path") + + @pytest.fixture(scope="module") def client(core_url, serving_url, allow_dirty): # Get client for core and serving @@ -45,6 +53,94 @@ def client(core_url, serving_url, allow_dirty): return client +def test_get_batch_features_with_file(client): + file_fs1 = FeatureSet( + "file_feature_set", + features=[Feature("feature_value", ValueType.STRING)], + entities=[Entity("entity_id", ValueType.INT64)], + max_age=Duration(seconds=100), + ) + + client.apply(file_fs1) + file_fs1 = client.get_feature_set(name="file_feature_set", version=1) + + N_ROWS = 10 + time_offset = datetime.utcnow().replace(tzinfo=pytz.utc) + features_1_df = pd.DataFrame( + { + "datetime": [time_offset] * N_ROWS, + "entity_id": [i for i in range(N_ROWS)], + "feature_value": [f"{i}" for i in range(N_ROWS)], + } + ) + client.ingest(file_fs1, features_1_df) + + # Rename column (datetime -> event_timestamp) + features_1_df = features_1_df.rename(columns={"datetime": "event_timestamp"}) + + to_avro(df=features_1_df, file_path_or_buffer="file_feature_set.avro") + + feature_retrieval_job = client.get_batch_features( + entity_rows="file://file_feature_set.avro", feature_ids=["file_feature_set:1:feature_value"] + ) + + output = feature_retrieval_job.to_dataframe() + print(output.head()) + + assert output["entity_id"].to_list() == [int(i) for i in output["file_feature_set_v1_feature_value"].to_list()] + + +def test_get_batch_features_with_gs_path(client, gcs_path): + gcs_fs1 = FeatureSet( + "gcs_feature_set", + features=[Feature("feature_value", ValueType.STRING)], + entities=[Entity("entity_id", ValueType.INT64)], + max_age=Duration(seconds=100), + ) + + client.apply(gcs_fs1) + gcs_fs1 = client.get_feature_set(name="gcs_feature_set", version=1) + + N_ROWS = 10 + time_offset = datetime.utcnow().replace(tzinfo=pytz.utc) + features_1_df = pd.DataFrame( + { + "datetime": [time_offset] * N_ROWS, + "entity_id": [i for i in range(N_ROWS)], + "feature_value": [f"{i}" for i in range(N_ROWS)], + } + ) + client.ingest(gcs_fs1, features_1_df) + + # Rename column (datetime -> event_timestamp) + features_1_df = features_1_df.rename(columns={"datetime": "event_timestamp"}) + + # Output file to local + file_name = "gcs_feature_set.avro" + to_avro(df=features_1_df, file_path_or_buffer=file_name) + + uri = urlparse(gcs_path) + bucket = uri.hostname + ts = int(time.time()) + remote_path = str(uri.path).strip("/") + f"{ts}/{file_name}" + + # Upload file to gcs + storage_client = storage.Client(project=None) + bucket = storage_client.get_bucket(bucket) + blob = bucket.blob(remote_path) + blob.upload_from_filename(file_name) + + feature_retrieval_job = client.get_batch_features( + entity_rows=f"{gcs_path}{ts}/*", + feature_ids=["gcs_feature_set:1:feature_value"] + ) + + output = feature_retrieval_job.to_dataframe() + print(output.head()) + + assert output["entity_id"].to_list() == [int(i) for i in output["gcs_feature_set_v1_feature_value"].to_list()] + + def test_order_by_creation_time(client): proc_time_fs = FeatureSet( "processing_time", diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index b37770a83f9..8ea472b6620 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -2,3 +2,4 @@ def pytest_addoption(parser): parser.addoption("--core_url", action="store", default="localhost:6565") parser.addoption("--serving_url", action="store", default="localhost:6566") parser.addoption("--allow_dirty", action="store", default="False") + parser.addoption("--gcs_path", action="store", default="gs://feast-templocation-kf-feast/") diff --git a/tests/e2e/requirements.txt b/tests/e2e/requirements.txt index 6b999421c04..0ba345a000f 100644 --- a/tests/e2e/requirements.txt +++ b/tests/e2e/requirements.txt @@ -1,6 +1,7 @@ mock==2.0.0 numpy==1.16.4 pandas==0.24.2 +pandavro==1.5.* pytest==5.2.1 pytest-benchmark==3.2.2 pytest-mock==1.10.4