pydataframer-databricks 0.1.0__py3-none-any.whl

pydataframer_databricks/__init__.py

@@ -0,0 +1,11 @@
+"""
+pydataframer-databricks: Databricks connector for Dataframer
+"""
+
+from .connectors import DatabricksConnector, FileType, DatasetType
+
+__all__ = [
+    "FileType",
+    "DatasetType",
+    "DatabricksConnector",
+]

pydataframer_databricks/connectors.py

@@ -0,0 +1,251 @@
+from enum import Enum
+
+
+class DatasetType(Enum):
+    """Dataset type enumeration matching Dataframer backend."""
+    SINGLE_FILE = "SINGLE_FILE"
+    MULTI_FILE = "MULTI_FILE"
+    MULTI_FOLDER = "MULTI_FOLDER"
+
+
+class FileType(Enum):
+    """File type enumeration matching Dataframer backend."""
+    MD = "md"
+    TXT = "txt"
+    CSV = "csv"
+    PDF = "pdf"
+    JSON = "json"
+    JSONL = "jsonl"
+     
+
+class DatabricksConnector:
+    """
+    Databricks connector for Dataframer workflows.
+    
+    This class provides methods to interact with Databricks SQL, fetch sample data,
+    and load generated data into Databricks tables.
+    
+    Parameters
+    ----------
+    dbutils : DBUtils
+        The dbutils object from your Databricks notebook context.
+        This is automatically available in Databricks notebooks.
+    
+    Examples
+    --------
+    >>> databricks_connector = DatabricksConnector(dbutils)
+    >>> df = databricks_connector.fetch_sample_data(
+    ...     num_items_to_select=25,
+    ...     table_name="samples.bakehouse.media_customer_reviews"
+    ... )
+    >>> df.head()
+    """
+    
+    def __init__(self, dbutils):
+        """
+        Initialize the Databricks connector.
+        
+        Parameters
+        ----------
+        dbutils : DBUtils
+            The dbutils object from your Databricks notebook context.
+        """
+        self.dbutils = dbutils
+    
+    def get_connection(self):
+        """
+        Return an authenticated Databricks SQL connection.
+        
+        Returns
+        -------
+        Connection
+            A Databricks SQL connection object.
+        """
+        from databricks import sql
+        from databricks.sdk.core import Config, oauth_service_principal
+
+        server_hostname = self.dbutils.secrets.get("dataframer", "DATABRICKS_SERVER_HOSTNAME")
+        http_path = self.dbutils.secrets.get("dataframer", "DATABRICKS_HTTP_PATH")
+
+        def credential_provider():
+            config = Config(
+                host=f"https://{server_hostname}",
+                client_id=self.dbutils.secrets.get("dataframer", "DATABRICKS_CLIENT_ID"),
+                client_secret=self.dbutils.secrets.get("dataframer", "DATABRICKS_CLIENT_SECRET"),
+            )
+            return oauth_service_principal(config)
+
+        return sql.connect(
+            server_hostname=server_hostname,
+            http_path=http_path,
+            credentials_provider=credential_provider,
+            user_agent_entry="dataframer_user_agent",
+        )
+    
+    def fetch_sample_data(self, num_items_to_select, table_name):
+        """
+        Fetch sample data from a Databricks table and return it as a Pandas DataFrame.
+        
+        Parameters
+        ----------
+        num_items_to_select : int
+            Number of rows to fetch from the table.
+        table_name : str
+            Fully qualified table name (e.g., "catalog.schema.table").
+        
+        Returns
+        -------
+        pd.DataFrame
+            A Pandas DataFrame containing the sample data.
+        
+        Examples
+        --------
+        >>> databricks_connector = DatabricksConnector(dbutils)
+        >>> df = databricks_connector.fetch_sample_data(
+        ...     num_items_to_select=25,
+        ...     table_name="samples.bakehouse.media_customer_reviews"
+        ... )
+        >>> df.head()
+        """
+        import pandas as pd
+
+        query = f"""
+            SELECT *
+            FROM {table_name}
+            LIMIT {num_items_to_select}
+        """
+
+        try:
+            with self.get_connection() as connection:
+                with connection.cursor() as cursor:
+                    cursor.execute(query)
+                    rows = cursor.fetchall()
+                    columns = [desc[0] for desc in cursor.description]
+        except Exception as e:
+            error_msg = f"Failed to fetch data from table `{table_name}`"
+            print(f"{error_msg}: {str(e)}")
+            print("Verify table exists, is accessible, and you have SELECT permissions")
+            raise RuntimeError(f"{error_msg}: {str(e)}") from e
+
+        return pd.DataFrame(rows, columns=columns)
+    
+    def load_generated_data(self, table_name, downloaded_zip, dataset_type, file_type):
+        """
+        Load generated samples from a ZIP file into a Databricks table using Databricks SQL.
+
+        Parameters
+        ----------
+        table_name : str
+            Target table name (e.g., "catalog.schema.table")
+        downloaded_zip : file-like
+            ZIP file object containing the generated data file
+        dataset_type : DatasetType
+            Type of dataset structure (DatasetType.SINGLE_FILE, DatasetType.MULTI_FILE, or DatasetType.MULTI_FOLDER)
+        file_type : FileType
+            Type of file in the ZIP (FileType.CSV, FileType.JSON, FileType.JSONL, etc.)
+        
+        Examples
+        --------
+        >>> databricks_connector = DatabricksConnector(dbutils)
+        >>> with open("samples.zip", "rb") as f:
+        ...     databricks_connector.load_generated_data(
+        ...         table_name="my_catalog.my_schema.my_table",
+        ...         downloaded_zip=f,
+        ...         dataset_type=DatasetType.SINGLE_FILE,
+        ...         file_type=FileType.CSV
+        ...     )
+        """
+        import zipfile
+        import pandas as pd
+        from io import BytesIO
+        
+        zip_buffer = BytesIO(downloaded_zip.read())
+
+        if dataset_type == DatasetType.SINGLE_FILE:
+            try:
+                with zipfile.ZipFile(zip_buffer) as z:
+                    file_list = z.namelist()
+                    
+                    generated_data_files = [f for f in file_list if f.lower().endswith(f'.{file_type.value}')]
+                    
+                    if len(generated_data_files) != 1:
+                        error_msg = f"Expected exactly one .{file_type.value} file in ZIP"
+                        print(f"{error_msg}. Available files: {file_list}")
+                        raise ValueError(error_msg)
+                    
+                    data_filename = generated_data_files[0]
+                    data_bytes = z.read(data_filename)
+                    print(f"Found {file_type.value} file: {data_filename}")
+            
+            except zipfile.BadZipFile as e:
+                error_msg = "Invalid or corrupted ZIP file"
+                print(f"{error_msg}: {str(e)}")
+                raise ValueError(f"{error_msg}: {str(e)}") from e
+            except ValueError:
+                raise
+            except Exception as e:
+                error_msg = "Failed to extract file from ZIP"
+                print(f"{error_msg}: {str(e)}")
+                raise RuntimeError(f"{error_msg}: {str(e)}") from e
+
+            if file_type == FileType.CSV:
+                pandas_df = pd.read_csv(BytesIO(data_bytes))
+            elif file_type == FileType.JSON:
+                # TODO: Implement JSON file handling
+                pass
+            elif file_type == FileType.JSONL:
+                # TODO: Implement JSONL file handling
+                pass
+            else:
+                raise ValueError(f"Unsupported file_type: {file_type}. Supported: CSV, JSON, JSONL for SINGLE_FILE datasets")
+            
+            with self.get_connection() as connection:
+                cursor = connection.cursor()
+
+                columns_sql = ", ".join(
+                    f"`{col}` STRING" for col in pandas_df.columns
+                )
+
+                try:
+                    cursor.execute(f"""
+                        CREATE OR REPLACE TABLE {table_name} (
+                            {columns_sql}
+                        )
+                    """)
+                except Exception as e:
+                    error_msg = f"Failed to create table `{table_name}`"
+                    print(f"{error_msg}: {str(e)}")
+                    print("Verify table name format (catalog.schema.table), permissions, and warehouse is running")
+                    cursor.close()
+                    raise RuntimeError(f"{error_msg}: {str(e)}") from e
+
+                insert_sql = f"""
+                    INSERT INTO {table_name}
+                    VALUES ({", ".join(["?"] * len(pandas_df.columns))})
+                """
+
+                try:
+                    cursor.executemany(
+                        insert_sql,
+                        pandas_df.values.tolist()
+                    )
+                except Exception as e:
+                    error_msg = f"Failed to insert data into table `{table_name}`"
+                    print(f"{error_msg}: {str(e)} | Rows attempted: {len(pandas_df)}")
+                    cursor.close()
+                    raise RuntimeError(f"{error_msg}: {str(e)}") from e
+
+                cursor.close()
+
+            print(f"✅ Table `{table_name}` saved successfully using Databricks SQL")
+        
+        elif dataset_type == DatasetType.MULTI_FILE:
+            # TODO: Implement MULTI_FILE handling
+            pass
+        
+        elif dataset_type == DatasetType.MULTI_FOLDER:
+            # TODO: Implement MULTI_FOLDER handling
+            pass
+        
+        else:
+            raise ValueError(f"Invalid dataset_type: {dataset_type}. Expected DatasetType enum")

pydataframer_databricks-0.1.0.dist-info/METADATA

@@ -0,0 +1,43 @@
+Metadata-Version: 2.4
+Name: pydataframer-databricks
+Version: 0.1.0
+Summary: Databricks connector for Dataframer
+Author-email: Dataframer <info@dataframer.ai>
+License: MIT
+Requires-Python: >=3.9
+Requires-Dist: databricks-sdk>=0.81.0
+Requires-Dist: databricks-sql-connector>=4.2.4
+Requires-Dist: pandas>=2.0.0
+Provides-Extra: dev
+Requires-Dist: pytest-cov>=4.1.0; extra == 'dev'
+Requires-Dist: pytest>=7.4.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# pydataframer-databricks
+
+Databricks connector package for Dataframer, providing seamless integration with Databricks SQL and data operations.
+
+## Installation
+
+```bash
+pip install pydataframer-databricks
+```
+
+## Building
+
+Requires [uv](https://docs.astral.sh/uv/) installed in your environment.
+
+```bash
+uv build
+```
+
+## Development
+
+```bash
+# Install with dev dependencies
+uv pip install -e ".[dev]"
+
+# Run tests
+pytest
+```
+

pydataframer_databricks-0.1.0.dist-info/RECORD

@@ -0,0 +1,5 @@
+pydataframer_databricks/__init__.py,sha256=piRrFtKpGGc2ctFDnMNblp5Whp6froRKXNeYkHnrw_o,214
+pydataframer_databricks/connectors.py,sha256=E4RlU30ADp0V27tuHWOai-7CM1YvmTInS_YonUpWMds,9191
+pydataframer_databricks-0.1.0.dist-info/METADATA,sha256=L-5a9ThsJYq_CTeAI22Zlo269NoXH3jRtyFOW_dPxAQ,891
+pydataframer_databricks-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+pydataframer_databricks-0.1.0.dist-info/RECORD,,

pydataframer_databricks-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any