databricks-ddbxutils 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

{databricks_ddbxutils-0.2.0.dist-info → databricks_ddbxutils-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: databricks-ddbxutils
-Version: 0.2.0
+Version: 0.4.0
 Summary: extends databricks dbutils
 Author: Haneul Kim
 Author-email: haneul.kim@data-dynamics.io
@@ -8,10 +8,14 @@ Requires-Python: >=3.11,<4.0
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: databricks-sdk (>=0.57.0,<0.58.0)
+Requires-Dist: cloudpickle (>=3.1.1,<4.0.0)
+Requires-Dist: databricks-sdk (>=0.64.0,<0.65.0)
 Requires-Dist: dotenv (>=0.9.9,<0.10.0)
 Requires-Dist: jinja2 (>=3.1.6,<4.0.0)
+Requires-Dist: pyarrow (==20.0.0)
+Requires-Dist: pyspark (>=4.0.0,<5.0.0)
 Requires-Dist: python-dateutil (>=2.9.0.post0,<3.0.0)
+Requires-Dist: pytz (>=2025.2,<2026.0)
 Description-Content-Type: text/markdown
 
 # databricks-ddbxutils

{databricks_ddbxutils-0.2.0.dist-info → databricks_ddbxutils-0.4.0.dist-info}/RECORD

@@ -1,9 +1,11 @@
 ddbxutils/__init__.py,sha256=ElK7lAZwcQh50qAVQ16t9Cfk0caZAYLJtKkDJDkLljk,158
+ddbxutils/datasources/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ddbxutils/datasources/pyfunc.py,sha256=AAPCdR9Pw7NDIbAAZi8yOB2T9usLTD2u9lpwqJMv2pI,4171
 ddbxutils/functions.py,sha256=eb1cM5tpJKB1MuzJ8ncTpGxTjZA70kDj_gJFQC6y7zg,1567
 ddbxutils/main.py,sha256=9pjd5VXAQOj-_lZsTojv-A4LUiZ-SLuK5oMgNwO4Fzc,488
 ddbxutils/widgets/__init__.py,sha256=aeFsmeixYUPaFFw4DqgvuQreBlkNwmBM4v2bCZUB2zU,1364
 ddbxutils/widgets/core.py,sha256=uBhbq5KKSyFT5CLRPsTw4zzeWYxGuJvImEQgh_iiEqk,989
-databricks_ddbxutils-0.2.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-databricks_ddbxutils-0.2.0.dist-info/METADATA,sha256=QQsJYSPRCg-fUKeCF05fCFLuozonYeTvX0KYyc0-XwM,2724
-databricks_ddbxutils-0.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-databricks_ddbxutils-0.2.0.dist-info/RECORD,,
+databricks_ddbxutils-0.4.0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+databricks_ddbxutils-0.4.0.dist-info/METADATA,sha256=7en42X_IwuMIDocSmWyIQfB_McsNIxFqk2oS6UpMzk4,2881
+databricks_ddbxutils-0.4.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+databricks_ddbxutils-0.4.0.dist-info/RECORD,,

ddbxutils/datasources/pyfunc.py

@@ -0,0 +1,122 @@
+import base64
+from dataclasses import dataclass
+
+import cloudpickle
+from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition
+from pyspark.sql.types import StructType, StructField, IntegerType
+
+
+@dataclass
+class PythonFunctionPartition(InputPartition):
+    """
+    Partition 정의: 각 파티션의 시작과 끝 범위
+    """
+    start: int
+    end: int
+
+
+class PythonFunctionReader(DataSourceReader):
+    """
+    DataSourceReader 구현
+    """
+
+    def __init__(self, schema: StructType, options: dict, serialized_func_b64):
+        self.schema = schema
+        self.options = options
+        self.func = cloudpickle.loads(base64.b64decode(serialized_func_b64))
+
+    def partitions(self):
+        lower = int(self.options.get('lowerLimit', '0'))
+        upper = int(self.options.get('upperLimit', '0'))
+        num_parts = int(self.options.get('numPartitions', '1'))
+        step = (upper - lower) // num_parts if num_parts > 0 else (upper - lower)
+        # print(f'step={step}')
+        parts = []
+        start = lower
+        for i in range(num_parts):
+            end = upper if i == num_parts - 1 else start + step
+            parts.append(PythonFunctionPartition(start, end))
+            start = end
+        return parts
+
+    def read(self, partition: PythonFunctionPartition):
+        for x in range(partition.start, partition.end):
+            # yield (self.func(x),)
+            yield self.func(x)
+
+
+class PythonFunctionDataSource(DataSource):
+    """
+    DataSource 구현
+
+    .. versionadded: 0.3.0
+
+    Notes
+    -----
+    user defined function 은 tuple, list, `pyspark.sql.types.Row`, `pyarrow.RecordBatch` 중에 하나의 type 을 가져야 합니다.
+
+    Examples
+    --------
+    >>> spark = ...
+
+    Use the default input partition implementation:
+
+    >>> def partitions(self):
+    ...     return [PythonFunctionPartition(1, 3)]
+
+    Subclass the input partition class:
+
+    >>> def partitions(self):
+    ...     return [PythonFunctionPartition(1, 3), PythonFunctionPartition(4, 6)]
+
+    Example in PySpark Shell 에서 `pyspark.sql.Row` 를 return 하는 함수를 사용하는 방법입니다.
+
+    >>> import base64
+    >>> import cloudpickle
+    >>> from ddbxutils.datasources.pyfunc import PythonFunctionDataSource
+    >>> spark.dataSource.register(PythonFunctionDataSource)
+    >>> # ...
+    >>> from pyspark.sql import Row
+    >>> def user_function_row(x) -> Row:
+    ...     from datetime import datetime
+    ...     from pytz import timezone
+    ...     return Row(str(x), str(x * x), datetime.now(timezone('Asia/Seoul')).strftime("%Y-%m-%d %H:%M:%S"))
+    ...
+    >>> df = (spark.read.format("pyfunc").
+    ...         schema('value1 string, value2 string, ts string').
+    ...         option("lowerLimit", "0").
+    ...         option("upperLimit", "10").
+    ...         option("numPartitions", "100").
+    ...         option("func", base64.b64encode(cloudpickle.dumps(user_function_row)).decode('utf-8')).
+    ...         load())
+    >>> df.show()
+
+    Example in PySpark Shell 에서 array 를 return 하는 함수를 사용하는 방법입니다.
+
+    >>> def user_function_row(x):
+    ...     from datetime import datetime
+    ...     from pytz import timezone
+    ...     return [str(x), str(x * x), datetime.now(timezone('Asia/Seoul')).strftime("%Y-%m-%d %H:%M:%S")]
+    ...
+    >>> df = (spark.read.format("pyfunc").
+    ...         schema('value1 string, value2 string, ts string').
+    ...         option("lowerLimit", "0").
+    ...         option("upperLimit", "10").
+    ...         option("numPartitions", "100").
+    ...         option("func", base64.b64encode(cloudpickle.dumps(user_function_row)).decode('utf-8')).
+    ...         load())
+    >>> df.show()
+    """
+
+    @classmethod
+    def name(cls):
+        return 'pyfunc'
+
+    def schema(self):
+        return StructType([StructField('value', IntegerType(), nullable=False)])
+
+    def reader(self, schema: StructType):
+        # options는 문자열이므로 필요시 변환하세요
+        # func = self.options.get('func', None)
+        func = self.options['func']
+        return PythonFunctionReader(self.schema(), self.options, func)