diff --git a/crates/pipeline-types/src/format/json.rs b/crates/pipeline-types/src/format/json.rs index cded48186f4..1a0086bfad1 100644 --- a/crates/pipeline-types/src/format/json.rs +++ b/crates/pipeline-types/src/format/json.rs @@ -140,6 +140,8 @@ pub enum JsonFlavor { /// JSON format accepted by the Kafka Connect `JsonConverter` class. #[serde(rename = "kafka_connect_json_converter")] KafkaConnectJsonConverter, + #[serde(rename = "pandas")] + Pandas, /// Parquet to-json format. /// (For internal use only) #[serde(skip)] diff --git a/crates/pipeline-types/src/serde_with_context/serde_config.rs b/crates/pipeline-types/src/serde_with_context/serde_config.rs index 4194a7308fe..51b6beeaf5e 100644 --- a/crates/pipeline-types/src/serde_with_context/serde_config.rs +++ b/crates/pipeline-types/src/serde_with_context/serde_config.rs @@ -140,6 +140,12 @@ impl From for SqlSerdeConfig { timestamp_format: TimestampFormat::String("%Y-%m-%dT%H:%M:%S%.f%:z"), decimal_format: DecimalFormat::String, }, + JsonFlavor::Pandas => Self { + time_format: TimeFormat::String("%H:%M:%S%.f"), + date_format: DateFormat::String("%Y-%m-%d"), + timestamp_format: TimestampFormat::MillisSinceEpoch, + decimal_format: DecimalFormat::String, + }, JsonFlavor::ParquetConverter => Self { time_format: TimeFormat::Nanos, date_format: DateFormat::String("%Y-%m-%d"), diff --git a/openapi.json b/openapi.json index 3df7598c0e5..1d3962822ed 100644 --- a/openapi.json +++ b/openapi.json @@ -3790,7 +3790,8 @@ "default", "debezium_mysql", "snowflake", - "kafka_connect_json_converter" + "kafka_connect_json_converter", + "pandas" ] }, "JsonParserConfig": { diff --git a/python/README.md b/python/README.md index 1f6582dc12f..9cf4e99147f 100644 --- a/python/README.md +++ b/python/README.md @@ -32,4 +32,12 @@ sphinx-apidoc -o . ../feldera make html ``` -To clean the build, run `make clean`. \ No newline at end of file +To clean the build, run `make clean`. + +## Testing + + +```bash +cd python +python3 -m unittest +``` diff --git a/python/feldera/_helpers.py b/python/feldera/_helpers.py index 87a6cdb6316..2da585e9f3b 100644 --- a/python/feldera/_helpers.py +++ b/python/feldera/_helpers.py @@ -32,3 +32,12 @@ def validate_connector_input_format(fmt: Format): if isinstance(fmt, JSONFormat) and fmt.config.get("update_format") is None: raise ValueError("update_format not set in the format config; consider using: .with_update_format()") + + +def chunk_dataframe(df, chunk_size=1000): + """ + Yield successive n-sized chunks from the given dataframe. + """ + + for i in range(0, len(df), chunk_size): + yield df.iloc[i:i + chunk_size] \ No newline at end of file diff --git a/python/feldera/output_handler.py b/python/feldera/output_handler.py index ea932624d6b..c265fc82509 100644 --- a/python/feldera/output_handler.py +++ b/python/feldera/output_handler.py @@ -20,7 +20,8 @@ def __init__(self, client: FelderaClient, pipeline_name: str, view_name: str, qu # the callback that is passed to the `CallbackRunner` def callback(df: pd.DataFrame, _: int): - self.buffer.append(df) + if not df.empty: + self.buffer.append(df) # sets up the callback runner self.handler = CallbackRunner(self.client, self.pipeline_name, self.view_name, callback, queue) @@ -38,4 +39,7 @@ def to_pandas(self): """ self.handler.join() - return pd.concat(self.buffer) + + if len(self.buffer) == 0: + return pd.DataFrame() + return pd.concat(self.buffer, ignore_index=True) diff --git a/python/feldera/rest/_httprequests.py b/python/feldera/rest/_httprequests.py index 4a0259020e1..76c7100eb18 100644 --- a/python/feldera/rest/_httprequests.py +++ b/python/feldera/rest/_httprequests.py @@ -9,6 +9,10 @@ from typing import Callable, Optional, Any, Union, Mapping, Sequence, List +def json_serialize(body: Any) -> str: + return json.dumps(body) if body else "" if body == "" else "null" + + class HttpRequests: def __init__(self, config: Config) -> None: self.config = config @@ -28,6 +32,7 @@ def send_request( content_type: str = "application/json", params: Optional[Mapping[str, Any]] = None, stream: bool = False, + serialize: bool = True, ) -> Any: """ :param http_method: The HTTP method to use. Takes the equivalent `requests.*` module. (Example: `requests.get`) @@ -36,6 +41,7 @@ def send_request( :param content_type: The value for `Content-Type` HTTP header. "application/json" by default. :param params: The query parameters part of this request. :param stream: True if the response is expected to be a HTTP stream. + :param serialize: True if the body needs to be serialized to JSON. """ self.headers["Content-Type"] = content_type @@ -71,7 +77,7 @@ def send_request( request_path, timeout=timeout, headers=headers, - data=json.dumps(body) if body else "" if body == "" else "null", + data=json_serialize(body) if serialize else body, params=params, stream=stream, ) @@ -102,8 +108,16 @@ def post( content_type: Optional[str] = "application/json", params: Optional[Mapping[str, Any]] = None, stream: bool = False, + serialize: bool = True, ) -> Any: - return self.send_request(requests.post, path, body, content_type, params, stream=stream) + return self.send_request( + requests.post, + path, + body, + content_type, + params, stream=stream, + serialize=serialize + ) def patch( self, diff --git a/python/feldera/rest/feldera_client.py b/python/feldera/rest/feldera_client.py index 6774361f93f..d8913e306fe 100644 --- a/python/feldera/rest/feldera_client.py +++ b/python/feldera/rest/feldera_client.py @@ -32,13 +32,19 @@ def __init__( """ :param url: The url to Feldera API (ex: https://try.feldera.com) :param api_key: The optional API key for Feldera - :param timeout: (optional) The amount of time in seconds that the cient will wait for a response beforing timing + :param timeout: (optional) The amount of time in seconds that the client will wait for a response before timing out. """ self.config = Config(url, api_key, timeout) self.http = HttpRequests(self.config) + try: + self.programs() + except Exception as e: + logging.error(f"Failed to connect to Feldera API: {e}") + raise e + def programs(self) -> list[Program]: """ Get all programs @@ -381,6 +387,8 @@ def push_to_pipeline( array: bool = False, force: bool = False, update_format: str = "raw", + json_flavor: str = None, + serialize: bool = True, ): """ Insert data into a pipeline @@ -394,8 +402,10 @@ def push_to_pipeline( :param force: If True, the data will be inserted even if the pipeline is paused :param update_format: JSON data change event format, used in conjunction with the "json" format, the default value is "insert_delete", other supported formats: "weighted", "debezium", "snowflake", "raw" - + :param json_flavor: JSON encoding used for individual table records, the default value is "default", other supported encodings: + "debezium_mysql", "snowflake", "kafka_connect_json_converter", "pandas" :param data: The data to insert + :param serialize: If True, the data will be serialized to JSON. True by default """ if format not in ["json", "csv"]: @@ -404,6 +414,9 @@ def push_to_pipeline( if update_format not in ["insert_delete", "weighted", "debezium", "snowflake", "raw"]: raise ValueError("update_format must be one of 'insert_delete', 'weighted', 'debezium', 'snowflake', 'raw'") + if json_flavor is not None and json_flavor not in ["default", "debezium_mysql", "snowflake", "kafka_connect_json_converter", "pandas"]: + raise ValueError("json_flavor must be one of 'default', 'debezium_mysql', 'snowflake', 'kafka_connect_json_converter', 'pandas'") + # python sends `True` which isn't accepted by the backend array = _prepare_boolean_input(array) force = _prepare_boolean_input(force) @@ -417,6 +430,9 @@ def push_to_pipeline( params["array"] = array params["update_format"] = update_format + if json_flavor is not None: + params["json_flavor"] = json_flavor + content_type = "application/json" if format == "csv": @@ -428,6 +444,7 @@ def push_to_pipeline( params=params, content_type=content_type, body=data, + serialize=serialize, ) def listen_to_pipeline( @@ -493,4 +510,4 @@ def listen_to_pipeline( if end and time.time() > end: break if chunk: - yield json.loads(chunk) \ No newline at end of file + yield json.loads(chunk) diff --git a/python/feldera/sql_context.py b/python/feldera/sql_context.py index 9a06519d1a2..462578bb26a 100644 --- a/python/feldera/sql_context.py +++ b/python/feldera/sql_context.py @@ -4,6 +4,7 @@ from typing import Optional, Dict, Callable +import pandas as pd from typing_extensions import Self from queue import Queue @@ -18,9 +19,9 @@ from feldera._callback_runner import CallbackRunner, _CallbackRunnerInstruction from feldera._helpers import ensure_dataframe_has_columns from feldera.formats import JSONFormat, CSVFormat, AvroFormat -from feldera._helpers import validate_connector_input_format from feldera.resources import Resources from feldera.enums import BuildMode, CompilationProfile +from feldera._helpers import validate_connector_input_format, chunk_dataframe def _table_name_from_sql(ddl: str) -> str: @@ -72,7 +73,7 @@ def __init__( # TODO: to be used for schema inference self.todo_tables: Dict[str, Optional[SQLTable]] = {} - self.http_input_buffer: list[Dict[str, dict | list[dict] | str]] = [] + self.http_input_buffer: list[Dict[str, pd.DataFrame]] = [] # buffer that stores all input connectors to be created # this is a Mapping[table_name -> list[Connector]] @@ -173,7 +174,16 @@ def __push_http_inputs(self): for input_buffer in self.http_input_buffer: for tbl_name, data in input_buffer.items(): - self.client.push_to_pipeline(self.pipeline_name, tbl_name, "json", data, array=True) + for datum in chunk_dataframe(data): + self.client.push_to_pipeline( + self.pipeline_name, + tbl_name, + "json", + datum.to_json(orient='records', date_format='epoch'), + json_flavor='pandas', + array=True, + serialize=False + ) self.http_input_buffer.clear() @@ -273,7 +283,7 @@ def connect_source_pandas(self, table_name: str, df: pandas.DataFrame): if tbl: # tbl.validate_schema(df) TODO: something like this would be nice - self.http_input_buffer.append({tbl.name: df.to_dict('records')}) + self.http_input_buffer.append({tbl.name: df}) return tbl = self.todo_tables.get(table_name) diff --git a/python/tests/test_pipeline.py b/python/tests/test_pipeline.py index ebd61d12245..ee13d208842 100644 --- a/python/tests/test_pipeline.py +++ b/python/tests/test_pipeline.py @@ -135,18 +135,18 @@ def test_listen_to_pipeline(self): name = str(uuid.uuid4()) self.test_create_pipeline(name, False) - TEST_CLIENT.start_pipeline(name) + TEST_CLIENT.pause_pipeline(name) t1 = threading.Thread(target=self.__listener, args=(name,)) t1.start() + TEST_CLIENT.start_pipeline(name) TEST_CLIENT.push_to_pipeline(name, "tbl", "csv", data) t1.join() assert self.result - TEST_CLIENT.pause_pipeline(name) TEST_CLIENT.shutdown_pipeline(name) TEST_CLIENT.delete_pipeline(name) diff --git a/python/tests/test_wireframes.py b/python/tests/test_wireframes.py index 7e3d1008de5..182c903fdd0 100644 --- a/python/tests/test_wireframes.py +++ b/python/tests/test_wireframes.py @@ -1,7 +1,7 @@ import time import unittest import pandas as pd -from kafka import KafkaProducer, KafkaConsumer, TopicPartition +from kafka import KafkaProducer, KafkaConsumer from kafka.admin import KafkaAdminClient, NewTopic from feldera import SQLContext, SQLSchema @@ -324,6 +324,31 @@ def test_pipeline_resource_config(self): assert TEST_CLIENT.get_pipeline(name).config["resources"] == config + def test_timestamp_pandas(self): + sql = SQLContext("test_timestamp_pandas", TEST_CLIENT).get_or_create() + + TBL_NAME = "items" + VIEW_NAME = "s" + + # backend doesn't support TIMESTAMP of format: "2024-06-06T18:06:28.443" + sql.register_table(TBL_NAME, SQLSchema({"id": "INT", "name": "STRING", "birthdate": "TIMESTAMP"})) + + sql.register_view(VIEW_NAME, f"SELECT * FROM {TBL_NAME}") + + df = pd.DataFrame({"id": [1, 2, 3], "name": ["a", "b", "c"], "birthdate": [ + pd.Timestamp.now(), pd.Timestamp.now(), pd.Timestamp.now() + ]}) + + sql.connect_source_pandas(TBL_NAME, df) + + out = sql.listen(VIEW_NAME) + + sql.run_to_completion() + + df = out.to_pandas() + + assert df.shape[0] == 3 + if __name__ == '__main__': unittest.main()